Catastrophic Forgetting › Preserving Pre-trained Capabilities
Language AI HandbookAll chapters Catastrophic Forgetting › Preserving Pre-trained CapabilitiesOne section · Sections in this chapter10 / 13 ↓
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
Beyond mitigating forgetting during training, you need strategies to actively preserve the valuable capabilities that pre-trained models bring. These strategies work by controlling the degree to which different parts of the model are allowed to change, exploiting the hierarchical structure of transformer architectures where different layers serve different functional roles.
The intuition for layer-wise strategies comes from a well-established finding in deep learning research: lower layers of a neural network tend to learn more general, transferable features, while higher layers learn more task-specific representations. In BERT-style models, the earliest transformer layers encode basic syntax and morphology, the middle layers capture semantic relationships and discourse structure, and the final layers specialize for the pre-training objective (masked language modeling or next sentence prediction). When fine-tuning, you generally want to allow the final layers to change substantially (they need to adapt to the new task) while protecting the early layers (they contain the general linguistic knowledge you want to preserve).
Layer-wise Learning Rates
Different layers encode different levels of abstraction. Earlier layers learn general features (syntax, basic semantics) while later layers specialize for the pre-training task. This suggests using different learning rates: small rates for early layers to limit how much they change, large rates for later layers to enable rapid adaptation to the new task.
def create_discriminative_learning_rates(model, base_lr=2e-5, decay_factor=0.9):
"""Create parameter groups with layer-wise learning rates."""
param_groups = []
# Get all transformer layers
num_layers = len(model.base_model.transformer.layer)
for i, layer in enumerate(model.base_model.transformer.layer):
# Earlier layers get smaller learning rates
layer_lr = base_lr * (decay_factor ** (num_layers - i - 1))
param_groups.append({"params": layer.parameters(), "lr": layer_lr})
# Classifier head gets full learning rate
param_groups.append(
{"params": model.classifier.parameters(), "lr": base_lr}
)
return param_groups
# Example usage
param_groups = create_discriminative_learning_rates(
model_standard, base_lr=2e-5
)Layer-wise learning rates: Encoder layer 0: 1.18e-05 Encoder layer 1: 1.31e-05 Encoder layer 2: 1.46e-05 Encoder layer 3: 1.62e-05 Encoder layer 4: 1.80e-05 Encoder layer 5: 2.00e-05 Classifier head: 2.00e-05
This discriminative schedule applies aggressive updates to the task-specific classifier head while protecting the foundational features in early encoder layers with much smaller learning rates. The decay factor of 0.9 means each layer going from the top down receives 90% of the learning rate of the layer above it, creating a smooth gradient from full adaptation at the output to near-stability at the input.
Gradual Unfreezing
Another approach is to initially freeze early layers and progressively unfreeze them during training. The core idea is that the early epochs are the most dangerous for forgetting: the model is adapting quickly to a new distribution, and large gradient updates can cause irreversible damage to general representations. By keeping those layers frozen during the most volatile phase and only thawing them once the model has found a stable trajectory for the new task, we can achieve most of the adaptation benefit while limiting the forgetting cost.
def freeze_layers(model, num_layers_to_freeze):
"""Freeze the first n transformer layers."""
for i, layer in enumerate(model.base_model.transformer.layer):
if i < num_layers_to_freeze:
for param in layer.parameters():
param.requires_grad = False
else:
for param in layer.parameters():
param.requires_grad = True
def count_trainable_params(model):
"""Count trainable parameters."""
return sum(p.numel() for p in model.parameters() if p.requires_grad)Loading weights: 0%| | 0/100 [00:00<?, ?it/s]
[transformers] [1mDistilBertForSequenceClassification LOAD REPORT[0m from: distilbert-base-uncased Key | Status | ------------------------+------------+- vocab_layer_norm.bias | UNEXPECTED | vocab_layer_norm.weight | UNEXPECTED | vocab_transform.bias | UNEXPECTED | vocab_transform.weight | UNEXPECTED | vocab_projector.bias | UNEXPECTED | classifier.bias | MISSING | classifier.weight | MISSING | pre_classifier.weight | MISSING | pre_classifier.bias | MISSING | Notes: - UNEXPECTED: can be ignored when loading from different task/architecture; not ok if you expect identical arch. - MISSING: those params were newly initialized because missing from the checkpoint. Consider training on your downstream task.
Gradual unfreezing schedule: Freeze 2 layers: 52,779,266 / 66,955,010 trainable (78.8%) Freeze 1 layers: 59,867,138 / 66,955,010 trainable (89.4%) Freeze 0 layers: 66,955,010 / 66,955,010 trainable (100.0%)
Starting with most parameters frozen limits the plasticity of the model, preventing large weight updates that could destroy pre-trained knowledge. As training stabilizes, more layers are unfrozen to allow fine-grained adaptation. The typical pattern is to begin training only the classifier head (which has no pre-trained knowledge to preserve, since it was randomly initialized), then progressively unfreeze layers from the top down, giving each phase enough epochs to stabilize before opening the next set of layers.
Evaluation-Driven Stopping
Perhaps the most reliable way to preserve capabilities is to continuously monitor them and stop training when degradation exceeds an acceptable threshold. This approach directly addresses the measurement gap that allows forgetting to go undetected: instead of waiting until deployment to discover capability losses, you detect them in real-time during training.
class ForgettingAwareTrainer:
"""Trainer that stops when forgetting exceeds threshold."""
def __init__(self, model, max_forgetting_pct=10.0):
self.model = model
self.max_forgetting_pct = max_forgetting_pct
self.baseline_scores = {}
def set_baseline(self, score_fn, capability_name):
"""Record baseline performance on capability."""
self.baseline_scores[capability_name] = score_fn(self.model)
def check_forgetting(self, score_fn, capability_name):
"""Check if forgetting exceeds threshold."""
if capability_name not in self.baseline_scores:
return False, 0.0
current = score_fn(self.model)
baseline = self.baseline_scores[capability_name]
# Calculate percentage degradation
if baseline != 0:
degradation = (baseline - current) / abs(baseline) * 100
else:
degradation = 0.0
return degradation > self.max_forgetting_pct, degradation# Demonstrate usage with the trainer
trainer = ForgettingAwareTrainer(model, max_forgetting_pct=10.0)
# Define a mock scoring function (e.g., accuracy on held-out task)
def get_capability_score(m):
return 0.85
# Set the baseline before training
trainer.set_baseline(get_capability_score, "reasoning")
# Check for forgetting (simulating a degradation)
def get_degraded_score(m):
return 0.75 # Performance dropped significantly
stop_signal, degradation = trainer.check_forgetting(
get_degraded_score, "reasoning"
)Degradation detected: 11.8% Stop training signal: True
The trainer detects that the performance drop (from 0.85 to 0.75) exceeds the allowable threshold, triggering a stop signal. This active monitoring acts as a safety brake. This keeps the model retains needed capabilities throughout the fine-tuning process. In a production fine-tuning pipeline, you would call check_forgetting at the end of each epoch with real scoring functions that probe the capabilities you care about, and halt training immediately if the stop signal fires.