Source-linked AI summary
An Empirical Study of Multi-Task Learning on BERT for Biomedical Text Mining
Yifan Peng, Qingyu Chen, Zhiyong Lu
TL;DR
Biomedical and clinical MTL evidence has mostly covered one or two tasks, leaving broader task combinations to examine. The paper evaluates shared-encoder BERT models across biomedical and clinical NLP tasks; MT-BERT reports gains in both domains and state-of-the-art average performance for the fine-tuned model.
Problem
Biomedical and clinical MTL has mostly been studied with one or two tasks, leaving its performance across broader task sets to examine.
Source: §5 · p. 7Method
The study compares independent BERT with MT-BERT refinement and task-specific fine-tuning across eight biomedical and clinical NLP tasks.
Source: §1 · p. 1Results
Takeaways & Limitations
Pairwise MTL shows which tasks can improve or decrease others, informing model choice for new problems with limited training resources.
Source: §1 · pp. 1–2MTL training across the BLUE benchmark is not always effective, and the data characteristics that predict success remain unclear.
Source: §5 · p. 8
Abstract
from arXiv · showhide
Multi-task learning (MTL) has achieved remarkable success in natural language processing applications. In this work, we study a multi-task learning model with multiple decoders on varieties of biomedical and clinical natural language processing tasks such as text similarity, relation extraction, named entity recognition, and text inference. Our empirical results demonstrate that the MTL fine-tuned models outperform state-of-the-art transformer models (e.g., BERT and its variants) by 2.0% and 1.3% in biomedical and clinical domains, respectively. Pairwise MTL further demonstrates more details about which tasks can improve or decrease others. This is particularly helpful in the context that researchers are in the hassle of choosing a suitable model for new problems. The code and models are publicly available at https://github.com/ncbi-nlp/bluebert
1 Introduction
The paper investigates whether BERT-based multi-task learning can improve biomedical and clinical NLP performance across diverse tasks. It reports gains on BLUE and examines how task pairings affect one another.
1 Introduction
- MTL shares representations across parallel tasks, which can increase training sample size and support generalization when labeled medical datasets are scarce.The paper frames this as useful for domains where labeled data are difficult to collect for deep learning.Source: §1 · p. 1
- Across biomedical and clinical BLUE evaluations, MTL improved performance by over 2% on average and achieved state-of-the-art results on four tasks.The authors also report state-of-the-art average performance for MT-BERT-Fine-Tune and continued benefit from MT-BERT-Refinement.Source: §1 · pp. 1–2
- The study compares independent BERT, MT-BERT-Refinement, and task-specific fine-tuning of the refined model across biomedical and clinical tasks.The evaluation uses BLUE, which spans text similarity, relation extraction, and named entity recognition among other tasks.Source: §1 · p. 1
- Pairwise MTL indicates which tasks can improve or decrease others, while the code and pretrained models are publicly available.The authors present these results as potentially useful when choosing a model for new problems with limited training resources.Source: §1 · pp. 1–2
2 Related work
Biomedical MTL research has often focused on one task across corpora or multiple tasks within one corpus. This study extends that scope to eight corpora and four task types, addressing mixed evidence about cross-task, cross-corpus MTL.
2 Related work
- Earlier NLP research examined features such as learning curves and entropy distributions as possible predictors of when MTL gains occur.The paper situates its study within broader work on conditions associated with MTL improvement.Source: §2 · p. 2
- Prior biomedical and clinical MTL has mainly used one task with multiple corpora or multiple tasks with annotations from a single corpus.Examples include NER across corpora and joint entity–relation extraction on individual datasets.Source: §2 · p. 2
- The study jointly learns across eight corpora and four task types, extending biomedical and clinical MTL beyond the one- or two-task settings common in prior work.The authors aim to clarify conditions and task relationships associated with MTL gains over single-task setups.Source: §2 · p. 2
- Prior work reports significant improvements in medicine tasks, but results for MTL across different tasks and corpora have been absent or mixed.This leaves the conditions under which cross-task, cross-corpus learning yields gains unresolved.Source: §2 · p. 2
3 Multi-task model
MT-BERT uses BERT layers shared across tasks and task-specific layers on top, with the shared representation trained through multi-task objectives.
3 Multi-task model
- Inputs may be a sentence or a packed sentence pair, with [CLS] and [SEP] tokens added and sequences truncated to the model’s maximum length.For sequence tagging, sentences longer than 30 words are split into shorter subsequences.Source: §3 · p. 2
- MT-BERT shares BERT layers that encode input tokens into contextual representations, then uses a fully connected task-specific layer for each task.Figure 1 presents the model architecture; the shared representation is trained with multi-task objectives.Source: §3 · pp. 2–3, Figure 1 · p. 3
3.1 Sentence similarity
Sentence similarity is modeled as a real-valued prediction from the [CLS] representation of the sentence pair, trained against the pair’s real-valued score.
3.1 Sentence similarity
- The sentence-pair similarity score is computed from the [CLS] embedding as sim(X1, X2) = ah0 + b and trained with mean squared error.The target y is the real-value similarity score for the sentence pair.Source: §3.1 · p. 3
3.2 Relation extraction
Relation extraction is framed as binary classification over a sentence after replacing the two arguments of interest with predefined tags. A fully connected layer and softmax predict the relation class, and training uses categorical crossentropy.
3.2 Relation extraction
- Relation extraction treats sentences containing two arguments of interest as binary relation-extraction examples.The arguments are replaced with predefined tags, such as GENE or DRUG, so the task can be treated as single-sentence classification.Source: §3.2 · p. 3
- A fully connected layer and softmax over the [CLS] embedding predict the relation class.The class probability is P(c|X) = softmax(ah0 + b).Source: §3.2 · p. 3
- Training uses categorical crossentropy, with the loss comparing predicted classes against the ground-truth class.The supplied loss expression defines the ground-truth indicator for each class.Source: §3.2 · p. 3
3.3 Inference
Inference packs a premise and hypothesis into one sequence and classifies their logical relation. The [CLS] embedding is used to predict that relation, with categorical crossentropy used for training.
3.3 Inference
- The inference task predicts the logical relation between a premise and a hypothesis.The premise and hypothesis are packed into one sequence, making the task a single-sentence classification problem.Source: §3.3 · p. 3
- The model uses the [CLS] embedding of the packed sequence to predict the relation with a softmax layer.The probability is specified as P(R|P ⊕ H) = softmax(ah0 + b).Source: §3.3 · p. 3
- The inference classifier is trained using categorical crossentropy.Source: §3.3 · p. 3
3.4 Named entity recognition
For named entity recognition, the model produces a feature vector for each input token and predicts a label sequence with a softmax output layer. Training uses categorical crossentropy.
3.4 Named entity recognition
- BERT produces a feature vector sequence matching the input sequence length for named entity recognition.Source: §3.4 · p. 3
- The MTL model predicts the label sequence using a softmax output layer.The layer scales outputs for labels in the label set.Source: §3.4 · p. 3
- The named entity recognition task is trained using categorical crossentropy.Source: §3.4 · p. 3
3.5 The training procedure
MT-BERT training has three stages: BERT pretraining, multi-task refinement, and task-specific fine-tuning. Refinement updates shared and task-specific parameters using batches drawn across datasets; fine-tuning continues training on each task.
3.5 The training procedure
- The training procedure comprises BERT pretraining, multi-task refinement, and task-specific fine-tuning.Source: §3.5 · pp. 3–4
3.5.1 Pretraining
- Pretraining uses the BERT masked language modeling technique with the base model and a maximum input length of 128 tokens.Source: §3.5.1 · p. 4
3.5.2 Refining via Multi-task learning
- Refinement initializes shared layers from pretrained BERT and task-specific parameters randomly, then updates the model using a task-specific objective for a sampled task mini-batch.Mini-batches from all datasets are merged, and mini-batch stochastic gradient descent is used.Source: §3.5.2 · p. 4
3.5.3 Fine-tuning MT-BERT
- Fine-tuning continues training all layers on each task, replacing the final softmax and linear layers and using a smaller learning rate.The passage conditions the transfer rationale on the dataset not being drastically different in context from the other datasets.Source: §3.5.3 · p. 4
4 Experiments
Experiments compare single-task BERT with multi-task refinement and fine-tuning on eight BLUE tasks. Fine-tuned MTL improves clinical and biomedical benchmark scores, while pairwise and variant analyses show that outcomes depend on task relationships and pretraining.
4 Experiments
- The evaluation compares task-specific BERT baselines with MT-BERT-Refinement and MT-BERT-Fine-Tune across eight BLUE benchmark tasks.The tasks span text similarity, relation extraction, named entity recognition, and inference.Source: §4 · p. 4
4.2 Training
- Training used BioBERT, BlueBERT, and ClinicalBERT with Adamax; task-specific fine-tuning ran for at most 10 epochs at learning rate e−5.Texts were truncated to 128 tokens; refinement training allowed up to 100 epochs, used 0.1 dropout, and clipped gradient norms to 1.Source: §4.2 · pp. 4–5
4.1 Datasets
- The eight datasets cover clinical and biomedical text, including sentence-pair tasks, relation extraction, and entity recognition.Clinical examples include ClinicalSTS, MedNLI, i2b2 2010, and ShARe/CLEF; biomedical examples include ChemProt, DDI, and BC5CDR. The study used the same train, development, and test sets for fair comparison.Source: §4.1 · p. 4, Table 1 · p. 5
4.3 Results
- 81.9% clinical and 83.6% biomedical BLUE benchmark scores; MT-BlueBERT-Fine-Tuneclinical improved over BlueBERTclinical by 1.3%, and MT-BlueBERT-Fine-Tunebiomedical improved over BlueBERTbiomedical by 2.0% absolute.The clinical and biomedical models set new state-of-the-art results on two tasks each; DDI had the largest biomedical task-level improvement at 8.1%.Source: Table 2 · p. 5, §4.3 · pp. 5–6, Table 3 · p. 5
4.4.1 Pairwise MTL
- Pairwise MTL effects varied: ShARe/CLEFE benefited from all three clinical partners, while DDI and ChemProt benefited and BC5CDR barely did.ClinicalSTS helped MedNLI; MedNLI helped other tasks except ClinicalSTS partially, and i2b2 2010 RE could help or harm depending on its partner. Figure 2 encodes improvement, decrease, and no effect with green, red, and grey directed edges.Source: Figure 2 · p. 6, §4.4.1 · pp. 6–7
4.4.2 MTL on BERT variants
- MT-BlueBERT-Fine-Tuneclinical led all clinical tasks; ClinicalBERT-based MTL fell ∼1.6%, while biomedical BioBERT- and BlueBERT-based MTL results were comparable.The authors note that limited details about ClinicalBERT pretraining warrant further investigation.Source: Table 4 · p. 6, Table 5 · p. 6, §4.4.2 · p. 7
4.4.3 Results on all BLUE tasks
- Across all BLUE tasks, MT-BioBERT-Fine-Tune had the best average performance, with MT-BlueBERT-Fine-Tunebiomedical close behind.The results again favored biomedical-literature pretraining for biomedical tasks and biomedical literature plus clinical notes for clinical tasks; the authors suggest separate models for different text genres may help.Source: Table 6 · p. 7, §4.4.3 · p. 7
5 Conclusions and future work
The study recommends combining MTL refinement with task-specific fine-tuning, reporting cross-domain gains while noting that MTL is not uniformly effective. It identifies task-relationship characterization and alternative sharing approaches as future work.
5 Conclusions and future work
- 2.0% and 1.3% improvements in biomedical and clinical domains, respectively, were achieved by MT-BERT learned and fine-tuned on a different domain.The authors recommend combining MTL refinement with task-specific fine-tuning based on the evaluation results.Source: §5 · p. 7
- MTL brought significant improvements in 4 tasks.Source: §5 · p. 8
- MTL across all BLUE tasks is not always effective, and which data or task characteristics predict success remains unclear.Source: §5 · p. 8
- The study shares an encoder across tasks with task-specific decoders; other MTL approaches remain to be investigated.Examples include fine-tuning only task-specific layers, soft parameter sharing, and knowledge distillation.Source: §5 · p. 8
- The authors hope the work informs generalizable NLP models and task relations that can lead to gains over single-task setups.Source: §5 · p. 8