Catastrophic Forgetting › Implementing L2-SP Regularization
Language AI HandbookAll chapters Catastrophic Forgetting › Implementing L2-SP RegularizationOne section · Sections in this chapter9 / 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
Let's implement L2-SP (L2 regularization toward Starting Point), a simple but effective forgetting mitigation technique. The implementation is straightforward: store the pre-trained weights before fine-tuning begins, then add a penalty term to the loss function at every training step that grows in proportion to how far the current weights have drifted from those stored values.
class L2SPRegularizer:
"""L2 regularization toward pre-trained weights."""
def __init__(self, model, alpha=0.01):
self.alpha = alpha
# Store pre-trained weights
self.pretrained_weights = {}
for name, param in model.named_parameters():
if param.requires_grad:
self.pretrained_weights[name] = param.data.clone()
def penalty(self, model):
"""Calculate L2-SP penalty."""
loss = 0.0
for name, param in model.named_parameters():
if name in self.pretrained_weights:
diff = param - self.pretrained_weights[name]
loss += (diff**2).sum()
return self.alpha * lossThe L2SPRegularizer class encapsulates the entire regularization mechanism. The __init__ method captures a frozen snapshot of all trainable parameters at initialization time. The penalty method, called once per training step, computes the total squared drift across all parameters and scales it by alpha. The gradient of this penalty with respect to each parameter is simply , which points back toward the pre-trained value whenever the parameter has drifted.
def train_with_l2sp(
model, dataloader, optimizer, regularizer, device, num_epochs=10
):
"""Train with L2-SP regularization."""
history = {"loss": [], "acc": [], "reg_loss": []}
for epoch in range(num_epochs):
model.train()
total_loss = 0
total_reg = 0
correct = 0
total = 0
for batch in dataloader:
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
labels = batch["labels"].to(device)
optimizer.zero_grad()
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask,
labels=labels,
)
task_loss = outputs.loss
reg_loss = regularizer.penalty(model)
loss = task_loss + reg_loss
loss.backward()
optimizer.step()
total_loss += task_loss.item()
total_reg += reg_loss.item()
preds = outputs.logits.argmax(dim=-1)
correct += (preds == labels).sum().item()
total += labels.size(0)
history["loss"].append(total_loss / len(dataloader))
history["acc"].append(correct / total)
history["reg_loss"].append(total_reg / len(dataloader))
return historyNotice that the training loop now computes two loss terms: the standard cross-entropy task_loss from the model's outputs, and the reg_loss from the regularizer. These are summed before calling loss.backward(), so the gradient signal that flows back through the network incorporates both the incentive to improve at the fine-tuning task and the penalty for drifting from pre-trained weights. The reg_loss is tracked separately in history so we can monitor how the regularization pressure evolves over training.
# Compare standard fine-tuning vs L2-SP
# Reset model for fair comparison
model_standard = AutoModelForSequenceClassification.from_pretrained(
model_name, num_labels=2
).to(device)
model_l2sp = AutoModelForSequenceClassification.from_pretrained(
model_name, num_labels=2
).to(device)
# Standard fine-tuning
optimizer_std = torch.optim.AdamW(model_standard.parameters(), lr=2e-4)
history_standard = {"loss": [], "acc": []}
for epoch in range(10):
model_standard.train()
total_loss = 0
correct = 0
total = 0
for batch in dataloader:
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
labels = batch["labels"].to(device)
optimizer_std.zero_grad()
outputs = model_standard(
input_ids=input_ids, attention_mask=attention_mask, labels=labels
)
outputs.loss.backward()
optimizer_std.step()
total_loss += outputs.loss.item()
preds = outputs.logits.argmax(dim=-1)
correct += (preds == labels).sum().item()
total += labels.size(0)
history_standard["loss"].append(total_loss / len(dataloader))
history_standard["acc"].append(correct / total)
# L2-SP fine-tuning
optimizer_l2sp = torch.optim.AdamW(model_l2sp.parameters(), lr=2e-4)
regularizer = L2SPRegularizer(model_l2sp, alpha=0.01)
history_l2sp = train_with_l2sp(
model_l2sp, dataloader, optimizer_l2sp, regularizer, device, num_epochs=10
)# Measure parameter drift for both models
def measure_parameter_drift(model, model_name_ref):
"""Calculate total parameter drift from pre-trained weights."""
ref_model = AutoModelForSequenceClassification.from_pretrained(
model_name_ref, num_labels=2
)
total_drift = 0
total_params = 0
for (name1, param1), (name2, param2) in zip(
model.named_parameters(), ref_model.named_parameters()
):
if "classifier" not in name1: # Exclude task-specific head
drift = ((param1.cpu() - param2) ** 2).sum().item()
total_drift += drift
total_params += param1.numel()
return np.sqrt(total_drift / total_params)
drift_standard = measure_parameter_drift(model_standard, model_name)
drift_l2sp = measure_parameter_drift(model_l2sp, model_name)
Standard fine-tuning drift: 0.0010 L2-SP fine-tuning drift: 0.0001 Drift reduction: 86.2%
The L2-SP regularizer successfully constrains parameter drift while maintaining task performance. In practice, you would tune the regularization strength based on validation performance on both the target task and held-out tasks measuring prior capabilities. Too small an provides negligible protection; too large an prevents the model from learning the new task effectively. A good starting range is , and you should use a validation set that includes both new-task performance and a held-out capability probe to select the optimal value.