[Paper Review] On the Stability of Fine-tuning BERT: Misconceptions, Explanations, and Strong Baselines

Marius Mosbach, Maksym Andriushchenko+1 more|arXiv (Cornell University)|Jun 8, 2020

Nuclear reactor physics and engineering41 references211 citations

TL;DR

The paper shows that fine-tuning instability of BERT-based models is mainly due to optimization difficulties (vanishing gradients) and generalization variance, not catastrophic forgetting or small data, and introduces a simple, strong baseline to greatly improve stability.

ABSTRACT

Fine-tuning pre-trained transformer-based language models such as BERT has become a common practice dominating leaderboards across various NLP benchmarks. Despite the strong empirical performance of fine-tuned models, fine-tuning is an unstable process: training the same model with multiple random seeds can result in a large variance of the task performance. Previous literature (Devlin et al., 2019; Lee et al., 2020; Dodge et al., 2020) identified two potential reasons for the observed instability: catastrophic forgetting and small size of the fine-tuning datasets. In this paper, we show that both hypotheses fail to explain the fine-tuning instability. We analyze BERT, RoBERTa, and ALBERT, fine-tuned on commonly used datasets from the GLUE benchmark, and show that the observed instability is caused by optimization difficulties that lead to vanishing gradients. Additionally, we show that the remaining variance of the downstream task performance can be attributed to differences in generalization where fine-tuned models with the same training loss exhibit noticeably different test performance. Based on our analysis, we present a simple but strong baseline that makes fine-tuning BERT-based models significantly more stable than the previously proposed approaches. Code to reproduce our results is available online: https://github.com/uds-lsv/bert-stable-fine-tuning.

Motivation & Objective

• Investigate why fine-tuning BERT-based models is unstable across seeds.
• Evaluate commonly cited hypotheses (catastrophic forgetting, small datasets) as causes of instability.
• Decompose instability into optimization and generalization components.
• Propose a simple, robust fine-tuning baseline that improves stability and performance.

Proposed method

• Analyze fine-tuning stability across BERT, RoBERTa, and ALBERT on GLUE tasks.
• Examine gradients to identify optimization issues causing failures.
• Assess impact of bias correction in ADAM and learning rate warmup.
• Evaluate effect of keeping training iterations high (longer training) on stability.
• Propose and validate a baseline fine-tuning setup with bias correction and extended training.

Experimental results

Research questions

• RQ1 What causes the observed instability during fine-tuning of BERT-based models?
• RQ2 Are catastrophic forgetting and small dataset size the primary culprits for instability?
• RQ3 How do optimization dynamics (e.g., gradient vanishing) and generalization contribute to instability?
• RQ4 Can a simple baseline improve fine-tuning stability across architectures and datasets?

Key findings

• Instability is better explained by optimization difficulties (vanishing gradients) and late-stage generalization variance rather than catastrophic forgetting or small data alone.
• Failed runs show vanishing gradients in bottom layers, while successful runs have stronger gradients throughout training.
• Bias correction in Adam and warmup-like effects substantially improve stability, especially for BERT and ALBERT; RoBERTa benefits but to a lesser extent.
• Increasing the number of training iterations and driving training loss toward near zero yields more consistent development performance.
• A simple baseline with AdamW, bias correction, a learning rate of 2e-5, and 20 epochs achieves markedly lower variability across seeds and competitive mean/max performance on RTE, MRPC, and CoLA.
• The findings generalize to RoBERTa and ALBERT beyond BERT.