BERT Fine-tuning › Catastrophic Forgetting
Language AI HandbookAll chapters BERT Fine-tuning › Catastrophic ForgettingOne section · Sections in this chapter7 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Catastrophic forgetting occurs when fine-tuning overwrites the general knowledge BERT learned during pre-training. The model becomes highly specialized for the fine-tuning task but loses its ability to generalize.
The phenomenon where a neural network, when trained on a new task, rapidly forgets previously learned information. In the context of BERT fine-tuning, this means losing pre-trained language understanding while adapting to a specific downstream task.
Signs of Catastrophic Forgetting
Several symptoms indicate catastrophic forgetting:
- Validation loss increases after initial decrease
- Out-of-domain performance drops on examples unlike the training data
- Model becomes overconfident on training-like examples but fails on variations
- Pre-training task performance degrades (masked language modeling accuracy drops)
Prevention Strategies
Several techniques mitigate catastrophic forgetting:
1. Use small learning rates: The most important factor. Rates of 2e-5 to 5e-5 allow gradual adaptation.
2. Train for few epochs: 2-4 epochs is typically sufficient. More epochs increase forgetting risk.
3. Early stopping: Monitor validation loss and stop when it starts increasing.
class EarlyStopping:
"""Stop training when validation loss stops improving."""
def __init__(self, patience: int = 3, min_delta: float = 0.0):
self.patience = patience
self.min_delta = min_delta
self.best_loss = float("inf")
self.counter = 0
self.best_weights = None
def __call__(self, val_loss, model):
if val_loss < self.best_loss - self.min_delta:
self.best_loss = val_loss
self.counter = 0
# Save best weights
self.best_weights = {
k: v.clone() for k, v in model.state_dict().items()
}
return False
else:
self.counter += 1
if self.counter >= self.patience:
# Restore best weights
model.load_state_dict(self.best_weights)
return True
return False4. Regularization: Weight decay penalizes large weights, keeping them from growing unbounded. The AdamW optimizer applies weight decay directly to the weights:
where:
- : the model parameters at step
- : the learning rate
- , : bias-corrected first and second moment estimates from Adam
- : the weight decay coefficient (0.01 is standard for BERT)
- : small constant for numerical stability
The update has two parts. The first term is the standard Adam update, moving parameters in the direction that reduces loss. The second term shrinks every weight toward zero by a fraction proportional to its current magnitude. This shrinkage keeps weights close to their initial values, which for BERT means close to the pre-trained weights. The net effect: the model can adapt to the task, but large deviations from pre-training are penalized.
5. Layer freezing: Freeze lower BERT layers that encode general features, only fine-tuning upper layers.
def freeze_bert_layers(model, num_layers_to_freeze: int):
"""
Freeze the embedding layer and first N transformer layers.
Later layers and the classifier remain trainable.
"""
# Freeze embeddings
for param in model.bert.embeddings.parameters():
param.requires_grad = False
# Freeze specified encoder layers
for i, layer in enumerate(model.bert.encoder.layer):
if i < num_layers_to_freeze:
for param in layer.parameters():
param.requires_grad = False# Demonstrate the effect of freezing
model_demo = BertForClassification(num_classes=2)
# Count parameters before freezing
total_before = sum(p.numel() for p in model_demo.parameters())
trainable_before = sum(
p.numel() for p in model_demo.parameters() if p.requires_grad
)
# Freeze 6 layers
layers_frozen = 6
freeze_bert_layers(model_demo, num_layers_to_freeze=layers_frozen)
trainable_after = sum(
p.numel() for p in model_demo.parameters() if p.requires_grad
)
reduction_pct = 100 * (1 - trainable_after / trainable_before)Layer Freezing Effect ---------------------------------------- Total parameters: 109,483,778 Trainable before freezing: 109,483,778 (100.0%) Trainable after freezing 6 layers: 43,119,362 (39.4%) Reduction in trainable parameters: 60.6%
Freezing the bottom 6 layers reduces trainable parameters by roughly half. This significantly decreases memory requirements and training time while preserving the general language understanding encoded in lower layers. The remaining trainable parameters in layers 6-11 and the classifier can still adapt to the task.
6. Mixout regularization: Randomly replace model weights with pre-trained weights during training, maintaining proximity to the original model.
Gradual Unfreezing
A sophisticated approach progressively unfreezes layers during training:
def gradual_unfreeze(model, epoch, total_epochs, num_layers=12):
"""
Progressively unfreeze layers as training progresses.
Start with only the classifier trainable, end with all layers trainable.
"""
# Calculate how many layers to unfreeze
layers_to_unfreeze = int((epoch / total_epochs) * num_layers)
# Freeze all BERT layers first
for param in model.bert.parameters():
param.requires_grad = False
# Unfreeze from top (layer 11) to bottom
for i in range(num_layers - 1, num_layers - 1 - layers_to_unfreeze, -1):
if i >= 0:
for param in model.bert.encoder.layer[i].parameters():
param.requires_grad = True
# Classifier is always trainable
for param in model.classifier.parameters():
param.requires_grad = True
return layers_to_unfreezemodel_gradual = BertForClassification(num_classes=2)
total_params = sum(p.numel() for p in model_gradual.parameters())
total_epochs = 10
schedule = []
for epoch in range(1, total_epochs + 1):
unfrozen = gradual_unfreeze(
model_gradual, epoch, total_epochs=total_epochs, num_layers=12
)
trainable = sum(
p.numel() for p in model_gradual.parameters() if p.requires_grad
)
layers_desc = f"Top {unfrozen}" if unfrozen > 0 else "None"
schedule.append(
(epoch, layers_desc, trainable, 100 * trainable / total_params)
)
The schedule starts by training only the classifier head in early epochs, letting it adapt to the task while BERT's weights remain frozen. As training progresses, more transformer layers are unfrozen from the top (layer 11) down. By epoch 10, most layers are trainable. This gradual approach reduces forgetting risk because the classifier establishes useful gradients before deeper layers begin updating.