Source-linked AI summary
BAM! Born-Again Multi-Task Networks for Natural Language Understanding
Kevin Clark, Minh-Thang Luong, Urvashi Khandelwal, Christopher D. Manning, Quoc V. Le
TL;DR
Multi-task models often struggle to match single-task models. BAM distills single-task teachers into a multi-task student and uses teacher annealing to transition from distillation to supervised learning. On BERT fine-tuned for GLUE, the method consistently outperforms standard single-task and multi-task training, though reported gains are generally small relative to BERT’s self-supervised pre-training gains.
Problem
Multi-task neural networks often perform worse than single-task counterparts, creating a challenge for effective joint learning across tasks.
Method
BAM distills single-task teacher predictions into a multi-task student and uses teacher annealing to shift training from teacher outputs toward gold labels.
Results
BAM consistently outperforms standard single-task and multi-task training on BERT fine-tuned across GLUE tasks.
Takeaways & Limitations
The method produces robust multi-task gains and may make multi-task learning more broadly useful within NLP.
Takeaways & Limitations
Except for closely related tasks with small datasets, BAM’s gains are small compared with those from self-supervised pre-training, and definitive comparisons require significance testing over multiple trials.
Abstract
from arXiv · showhide
It can be challenging to train multi-task neural networks that outperform or even match their single-task counterparts. To help address this, we propose using knowledge distillation where single-task models teach a multi-task model. We enhance this training with teacher annealing, a novel method that gradually transitions the model from distillation to supervised learning, helping the multi-task model surpass its single-task teachers. We evaluate our approach by multi-task fine-tuning BERT on the GLUE benchmark. Our method consistently improves over standard single-task and multi-task training.
1 Introduction
BAM addresses the difficulty of making multi-task models match single-task performance by distilling knowledge from single-task teachers into a multi-task student. Teacher annealing then shifts training from teacher imitation toward gold labels, and the approach consistently outperforms standard single-task and multi-task training.
- Multi-task NLP models often perform worse than single-task counterparts, motivating knowledge distillation from single-task teachers to a multi-task model.
- BAM extends born-again networks to multi-task learning and compares Single→Multi distillation with alternative distillation strategies.
- Teacher annealing gradually transitions the student from teacher predictions to gold labels, helping it surpass its teachers.The teacher signal is richest early in training, while gold labels dominate later.
- BAM fine-tunes BERT across GLUE tasks and consistently outperforms standard single-task and multi-task training.
- Further analysis attributes multi-task benefits to better regularization and transfer between related tasks.
2 Related Work
Prior NLP work improves multi-task learning through architecture design, task selection, transfer, and distillation, whereas BAM focuses on the multi-task training algorithm.
- NLP multi-task research has emphasized selective parameter sharing and linguistically motivated task hierarchies.
- BAM differs by focusing on the multi-task training algorithm rather than neural architecture design or task-specific components.
- Knowledge distillation has compressed large models and ensembles, regularized multi-task agents, and transferred translation systems into multilingual models.
- Recent BERT multi-task studies use standard transfer or multi-task learning while exploring task pairs and task-specific components.
3 Methods
The method builds shared BERT-based multi-task models and trains them with supervised losses or teacher predictions. Single-task teachers distill into a multi-task student, while teacher annealing gradually restores reliance on gold labels.
- Multi-Task Setup: BERT produces contextual token representations, and multi-task models share nearly all parameters while retaining task-specific classifiers.Less than 0.01% of parameters are task-specific.
- Multi-Task Setup: Multi-task training shuffles examples from different tasks and minimizes the summed loss across tasks.
- Knowledge Distillation: Supervised learning minimizes loss against training labels, while knowledge distillation instead matches a teacher model’s predictions.
- Knowledge Distillation: Born-again distillation uses student and teacher models with the same architecture, with L2 loss for regression tasks.
- Multi-Task Distillation: Single-task models teach a shared multi-task model, with additional distillation strategies explored.
- Teacher Annealing: Teacher annealing mixes teacher predictions with gold labels, increasing λ linearly from 0 to 1 throughout training.Early training emphasizes distillation; late training emphasizes gold-standard labels so the student can surpass its teachers.
4 Experiments
The experiments evaluate BERT-based multi-task fine-tuning on nine English GLUE tasks, using task-balanced sampling, layerwise learning rates, and repeated trials to assess performance reliably.
- The GLUE benchmark contains 9 English natural language understanding tasks spanning entailment, paraphrase, similarity, sentiment, acceptability, and Winograd Schema.
- Task sampling uses probabilities proportional to |Dτ|^0.75, preventing very large datasets from overly dominating multi-task training.
- Layerwise learning rates decrease toward the input, preserving general features in pre-trained layers during fine-tuning.
- Multi-task models train for 6 epochs with batch size 128, α = 0.9, learning rate 1e-4, and BERT-Large pre-trained weights.
- Results report median scores from at least 20 random-seed trials, excluding WNLI from the average and using task-specific metrics.
5 Results
Across GLUE experiments, Single→Multi distillation with teacher annealing generally outperforms standard single-task and multi-task training while retaining one unified model. The results also show robustness, benefits from auxiliary tasks, and caution against overinterpreting single-run test scores.
- Single→Multi distillation improves or matches other methods on every GLUE task except STS, the benchmark’s only regression task.
- Standard multi-task training improves RTE but shows no improvement on the other tasks, whereas Single→Multi improves or matches performance broadly.
- Single→Multi gains exceed Single→Single gains, while additional distillation rounds provide no statistically significant improvement over Single→Multi.
- The method consistently outperforms standard single-task and multi-task training, with multi-task models showing more dependable performance across trials.
- Test-set results outperform or match published non-ensemble results, but individual-run scores and dev-set model selection limit definitive comparisons.
- Single-task fine-tuning improves results, but adds only a small gain over Single→Multi while producing task-specific rather than unified models.
- Teacher annealing is necessary: pure distillation with λ = 0 performs no better than standard multi-task learning.
- Auxiliary tasks improve RTE, with both closely related MNLI and relatively unrelated QQP, CoLA, and SST contributing benefits.
6 Discussion and Conclusion
Single→Multi distillation with teacher annealing consistently outperforms standard single-task and multi-task training, but its gains are generally small relative to BERT’s self-supervised transfer gains.
- Single→Multi distillation combined with teacher annealing consistently outperforms standard single-task or multi-task training.
- The method aims to make multi-task learning more broadly useful within NLP because robust multi-task gains have previously remained elusive.
- Except for closely related tasks with small datasets, the overall multi-task gains are small compared with transfer learning from BERT.