Source-linked AI summary
Visualizing and Understanding the Effectiveness of BERT
Yaru Hao, Li Dong, Furu Wei, Ke Xu
TL;DR
The paper addresses why pre-training followed by fine-tuning improves BERT's downstream trainability and generalization. It visualizes loss landscapes and optimization trajectories, finding wider optima, easier optimization, and more invariant lower layers that support transferable language representations.
Problem
It is unclear why the pre-training-then-fine-tuning paradigm improves performance and generalization across downstream tasks.
Method
The paper visualizes one-dimensional and two-dimensional loss landscapes and optimization trajectories for fine-tuning BERT, comparing them with random initialization.
Results
Pre-training provides good downstream initialization, wider optima, easier optimization, robustness to overfitting, and better generalization; lower layers are more invariant and transferable across tasks.
Takeaways & Limitations
Flat, wide optima and consistency between training-loss and generalization-error surfaces help explain BERT's generalization, while lower layers learn more transferable language representations.
Takeaways & Limitations
The experiments and conclusions are derived from BERT fine-tuning; extending the analysis to multi-task training and other pre-trained models remains future work.
Abstract
from arXiv · showhide
Language model pre-training, such as BERT, has achieved remarkable results in many NLP tasks. However, it is unclear why the pre-training-then-fine-tuning paradigm can improve performance and generalization capability across different tasks. In this paper, we propose to visualize loss landscapes and optimization trajectories of fine-tuning BERT on specific datasets. First, we find that pre-training reaches a good initial point across downstream tasks, which leads to wider optima and easier optimization compared with training from scratch. We also demonstrate that the fine-tuning procedure is robust to overfitting, even though BERT is highly over-parameterized for downstream tasks. Second, the visualization results indicate that fine-tuning BERT tends to generalize better because of the flat and wide optima, and the consistency between the training loss surface and the generalization error surface. Third, the lower layers of BERT are more invariant during fine-tuning, which suggests that the layers that are close to input learn more transferable representations of language.
1 Introduction
The paper investigates why BERT pre-training helps downstream trainability and generalization by visualizing fine-tuning loss landscapes and optimization trajectories. It reports wider optima, easier optimization, robustness to overfitting, and more transferable lower-layer representations.
- The paper asks why pre-training improves downstream trainability and generalization, using BERT as an illustrative case.
- The authors visualize one-dimensional and two-dimensional loss landscapes and optimization trajectories during fine-tuning.The one-dimensional curve compares movement from initialization to the fine-tuned solution, while the two-dimensional surface adds another optimization direction.
- Pre-training provides a good initial point across downstream tasks, producing wider optima and easier, faster optimization than random initialization.
- Fine-tuning BERT is robust to overfitting despite the model being highly over-parameterized for downstream tasks.
- Lower BERT layers are more invariant across tasks than higher layers, suggesting they learn transferable language representations.The paper evaluates this by visualizing loss landscapes for different layer groups.
2 Background: BERT
BERT is a pre-trained Transformer encoder that represents contextual information through stacked self-attention blocks. For downstream tasks, task-specific layers are added and the model parameters are fine-tuned.
- BERT is pre-trained on a large corpus with masked language modeling and next-sentence prediction objectives.
- For downstream tasks, task-specific layers are added and all BERT parameters are fine-tuned.
- BERT uses a multi-layer Transformer with self-attention blocks to encode contextual information.
- Each Transformer block transforms the preceding hidden states, producing final token representations hL_i.
3 Methodology
The paper visualizes BERT fine-tuning through one- and two-dimensional loss landscapes and projected optimization trajectories, comparing pretrained and randomly initialized models. These views characterize optimization directions, parameter-space surfaces, and training trajectories.
- Visualization framework: Three visualizations compare fine-tuning BERT with training from random initialization: 1D loss curves, 2D loss surfaces, and projected optimization trajectories.The same visualization algorithms can be applied to both learning paradigms.
- 3.1 One-dimensional Loss Curve: The 1D loss curve measures loss values along the optimization direction from initialized parameters θ0 to final parameters θ1.It is plotted by linear interpolation between θ0 and θ1.
- Parameter-space scope: Task-specific layers are excluded from the visualized parameter space by considering only BERT encoder parameters and fixing added layers to learned values.This isolates changes in the original BERT parameters during fine-tuning.
- 3.2 Two-dimensional Loss Surface: The 2D loss surface uses the target-dataset update direction δ1 = θ1 − θ0 and a second direction δ2 = θ2 − θ0 from another dataset.The two directions are normalized for visualization, and experiments report them as divergent and orthogonal in the high-dimensional parameter space.
- 3.3 Optimization Trajectory: The optimization trajectory projects each epoch’s parameter direction onto the 2D loss surface using deviations relative to the final direction and the visualization axes.At epoch i, θi defines δi = θi − θ0; cosine similarity with δ1 indicates the angle to the final optimization direction.
4 Experimental Setup
Experiments evaluate BERT-large and a matched randomly initialized architecture on four NLP datasets using accuracy. Fine-tuning follows standard BERT settings, while scratch training uses longer validated schedules to reach convergence.
- Datasets and evaluation: Experiments cover MNLI, RTE, SST-2, and MRPC, using the same data split as prior work and accuracy for evaluation.The datasets span natural language inference, textual entailment, sentiment classification, and paraphrase detection.
- BERT fine-tuning: The experiments use the pretrained cased BERT-large model with BERT-recommended hyperparameters and Adam optimization.Fine-tuning lasts 3, 4, or 5 epochs selected from the specified set.
- Hyperparameters: RTE and MRPC use batch size 32 with learning rate 1e-5, while MNLI and SST-2 use batch size 64 with learning rate 3e-5.These settings distinguish the optimization configurations across dataset groups.
- Training from scratch: Training from scratch keeps the BERT architecture and most hyperparameters but randomly initializes parameters and uses longer schedules for convergence.Scratch training uses 8 epochs for SST-2 and 16 epochs for the other datasets, validated on development data.
5 Pre-training Gets a Good Initial Point Across Downstream Tasks
Across four downstream datasets, BERT pre-training provides a better starting point than random initialization, producing wider optima, smoother optimization, faster convergence, and robustness to overfitting.
- 5.1 Pre-training Leads to Wider Optima: Fine-tuning BERT starts from a good initial point across datasets and reaches wider optima than training from scratch.The comparison uses four datasets and contrasts pre-trained with randomly initialized models.
- 5.2 Pre-training Eases Optimization on Downstream Tasks: Fine-tuning BERT follows a smoother path with generally decreasing training loss, whereas random initialization produces rougher paths requiring more carefully tuned optimization.The smoother path eases optimization and accelerates convergence.
- 5.2 Pre-training Eases Optimization on Downstream Tasks: Training from scratch requires more iterations to converge and often ends with higher training loss than fine-tuning BERT, including difficulty reducing loss early on RTE.These learning-curve observations hold across the four evaluated datasets.
- 5.2 Pre-training Eases Optimization on Downstream Tasks: Optimization from scratch can initially diverge from its final direction and encounter rough obstacles, while fine-tuning BERT finds a more direct, unimpeded route.The trajectory comparison includes MRPC and MNLI.
- 5.3 Pre-training-then-fine-tuning is Robust to Overfitting: Fine-tuning BERT remains robust to overfitting despite BERT-large having 345M parameters and being over-parameterized for the target datasets.Generalization error on the development set does not dramatically increase with additional training epochs.
6 Pre-training Helps to Generalize Better
Visualization results associate fine-tuning BERT with flatter, wider optima and better alignment between training loss and generalization error, especially relative to training from scratch on smaller datasets.
- 6.1 Wider Optima: Fine-tuned BERT reaches flatter and wider optima than models trained from scratch, a pattern associated with better generalization on unseen data.The one-dimensional curves use normalized axes to compare flatness.
- 6.2 Consistency Between Training Loss Surface and Generalization Error Surface: Fine-tuned BERT endpoints fall in wide regions with smaller generalization error, while scratch-trained models tend to occupy thinner optimum areas.The generalization measure is classification error on the development set.
- 6.2 Consistency Between Training Loss Surface and Generalization Error Surface: Training-loss and development-set error surfaces are consistent across datasets, so smaller training loss tends to accompany lower development-set error.Fine-tuned models also remain approximately optimal under subtle parameter perturbations.
- 6.2 Consistency Between Training Loss Surface and Generalization Error Surface: Scratch training generalizes more consistently on larger datasets such as MNLI and SST-2 than on smaller datasets such as MRPC and RTE.The results suggest scratch training usually requires more examples to generalize better than fine-tuning BERT.
7 Lower Layers of BERT are More Invariant and Transferable
Layer-wise analyses indicate that lower BERT layers are more invariant across tasks, while higher layers are more important for task-specific fine-tuning.
- Layer-wise loss surfaces: Lower layers have wider local optima, and rolling them back to pre-trained values does not dramatically hurt performance.The analysis groups layers 0th-7th as low layers, 8th-15th as middle layers, and 16th-23rd as high layers.
- Layer-wise loss surfaces: Rollback of high layers moves the model into a high-loss region, indicating that their optimization is critical to fine-tuning.
- Rollback evaluation: Development-set accuracy likewise changes little after rolling back low or middle layers, whereas rolling back high layers hurts performance.Table 1 reports accuracy for fine-tuned models and models with different layer groups rolled back.
- Rollback evaluation: Rollback of low layers even improves generalization on MNLI, supporting their greater invariance and transferability across tasks.The paper attributes this to lower layers learning more transferable language representations and higher layers learning more task-specific information.
8 Related Work
Related work examines pretrained models through linguistic probing, layer-wise transferability, and loss-landscape geometry, motivating this paper’s visualization-based analysis of BERT.
- Pretrained representations: Pretrained contextualized representations, including BERT, have shown strong performance across varied natural language processing tasks.
- Model analysis: Prior studies analyze pretrained models’ syntactic and semantic information using probing tasks and direct assessments of BERT’s syntactic abilities.
- Layer-wise representations: Earlier work reports that BERT layers differ in transferability, with lower layers encoding local syntax and higher layers capturing more complex semantics.
- Loss-landscape geometry: Computer-vision research connects loss-function geometry with generalization and characterizes flatness through connected regions around minima or Hessian eigenvalues.
- Loss-landscape geometry: Optimization methods have used the flatness assumption to seek broader optima associated with better generalization than standard SGD.
9 Conclusion
The paper uses loss-landscape and optimization-trajectory visualizations to explain BERT fine-tuning, finding wider optima, easier optimization, robust overfitting behavior, and transferable lower layers.
- Conclusion: Visualizations show that pre-training leads to wider optima and eases optimization compared with training from scratch.
- Conclusion: The pre-training-then-fine-tuning paradigm is robust to overfitting, despite the experiments focusing on BERT fine-tuning.
- Conclusion: Consistency between training-loss surfaces and generalization-error surfaces provides evidence relevant to BERT’s generalization capability.
- Conclusion: Lower BERT layers are more invariant and transferable across tasks.