Regularization Methods › Code Implementation
Language AI HandbookAll chapters Regularization Methods › Code ImplementationOne section · Sections in this chapter7 / 10 ↓
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 both EWC and SI from scratch using PyTorch and evaluate their ability to prevent catastrophic forgetting on a two-task sequential learning scenario.
Setup and Data Generation
We begin by importing the necessary libraries and creating a simple two-task dataset. We will use a multi-layer perceptron trained first on a binary classification task, then adapted to a second task.
import numpy as np
import torch
from torch.utils.data import DataLoader, TensorDataset
torch.manual_seed(42)
np.random.seed(42)
def make_task_data(n=500, task="A", noise=0.15):
X = np.random.randn(n, 2).astype(np.float32)
if task == "A":
y = ((X[:, 0] - X[:, 1]) < 0).astype(np.float32)
else:
y = ((X[:, 0] ** 2 + X[:, 1] ** 2) > 1.0).astype(np.float32)
flip = np.random.rand(n) < noise
y[flip] = 1 - y[flip]
return X, y
X_A, y_A = make_task_data(500, "A")
X_B, y_B = make_task_data(500, "B")
X_A_test, y_A_test = make_task_data(200, "A")
X_B_test, y_B_test = make_task_data(200, "B")
loader_A = DataLoader(
TensorDataset(torch.tensor(X_A), torch.tensor(y_A)),
batch_size=32,
shuffle=True,
)
loader_B = DataLoader(
TensorDataset(torch.tensor(X_B), torch.tensor(y_B)),
batch_size=32,
shuffle=True,
)Task A classifies points based on a linear boundary (sign of ), while task B uses a radial boundary (distance from the origin). These two tasks require different decision surfaces, so the network must learn different computations for each, making them a good stress test for catastrophic forgetting protection. The linear boundary for task A can be solved by a single linear layer; the radial boundary for task B requires nonlinear interactions between the two input features. A network that finds its weights adjusted to solve task A will generally need to reconfigure its internal representations to solve task B, creating conditions for forgetting.
Neural Network Architecture
We define a simple two-layer MLP. Keeping the architecture small makes the catastrophic forgetting effect visible and the importance weights interpretable.
class MLP(nn.Module):
def __init__(self, input_dim=2, hidden_dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, 1),
)
def forward(self, x):
return self.net(x).squeeze(-1)
def evaluate(model, X, y):
model.eval()
with torch.no_grad():
logits = model(torch.tensor(X))
preds = (torch.sigmoid(logits) > 0.5).float()
acc = (preds == torch.tensor(y)).float().mean().item()
return accThe MLP has two hidden layers with 64 units each, separated by ReLU activations. The output is a single logit, and we use binary cross-entropy with logits as the loss function. This architecture is deliberately simple so that the importance weight distributions and forgetting behavior are easy to observe without being obscured by the complexities of larger networks.
Baseline: Training Without Regularization
First, let's see the extent of catastrophic forgetting when we train sequentially without any protection.
def train_vanilla(model, loader, epochs=10, lr=1e-3):
optimizer = optim.Adam(model.parameters(), lr=lr)
criterion = nn.BCEWithLogitsLoss()
model.train()
for _ in range(epochs):
for X_batch, y_batch in loader:
optimizer.zero_grad()
loss = criterion(model(X_batch), y_batch)
loss.backward()
optimizer.step()
baseline_model = MLP()
train_vanilla(baseline_model, loader_A, epochs=15)
acc_A_after_A = evaluate(baseline_model, X_A_test, y_A_test)
train_vanilla(baseline_model, loader_B, epochs=15)
acc_A_after_B = evaluate(baseline_model, X_A_test, y_A_test)
acc_B_after_B = evaluate(baseline_model, X_B_test, y_B_test)=== Baseline (No Regularization) === Task A accuracy after training on A: 0.795 Task A accuracy after training on B: 0.470 Task B accuracy after training on B: 0.795 Forgetting on Task A: 0.325
The drop in task A accuracy after training on task B is catastrophic forgetting in action. The model learned task B effectively but lost much of its task A knowledge in the process. This is not a subtle degradation; the drop is substantial, corresponding to a model that has largely forgotten the linear boundary it learned. From the optimizer's perspective, this is entirely rational: it minimized the loss on task B as effectively as possible. It simply had no incentive to preserve task A.
EWC Implementation
Now let's implement EWC. The key component is computing the diagonal Fisher information matrix after training on task A.
def compute_fisher(model, loader):
model.eval()
fisher = {n: torch.zeros_like(p) for n, p in model.named_parameters()}
criterion = nn.BCEWithLogitsLoss()
n_samples = 0
for X_batch, y_batch in loader:
model.zero_grad()
logits = model(X_batch)
loss = criterion(logits, y_batch)
loss.backward()
for n, p in model.named_parameters():
if p.grad is not None:
fisher[n] += p.grad.data**2 * len(X_batch)
n_samples += len(X_batch)
for n in fisher:
fisher[n] /= n_samples
return fisher
class EWCTrainer:
def __init__(self, model, lambda_ewc=100.0):
self.model = model
self.lambda_ewc = lambda_ewc
self.fisher = {}
self.params_star = {}
def register_task(self, loader):
self.fisher = compute_fisher(self.model, loader)
self.params_star = {
n: p.data.clone() for n, p in self.model.named_parameters()
}
def ewc_loss(self):
loss = 0.0
for n, p in self.model.named_parameters():
if n in self.fisher:
loss += (self.fisher[n] * (p - self.params_star[n]) ** 2).sum()
return (self.lambda_ewc / 2) * loss
def train_with_ewc(self, loader, epochs=10, lr=1e-3):
optimizer = optim.Adam(self.model.parameters(), lr=lr)
criterion = nn.BCEWithLogitsLoss()
self.model.train()
for _ in range(epochs):
for X_batch, y_batch in loader:
optimizer.zero_grad()
task_loss = criterion(self.model(X_batch), y_batch)
reg_loss = self.ewc_loss()
total_loss = task_loss + reg_loss
total_loss.backward()
optimizer.step()The compute_fisher function accumulates squared gradients over the dataset, then normalizes by the number of samples. This produces the empirical diagonal Fisher. The EWCTrainer class wraps the model and manages the Fisher computation and importance-weighted regularization. The register_task method must be called after training on a task but before training on the next one; it snapshots both the importance weights and the parameter values to serve as the regularization anchors.
ewc_model = MLP()
ewc_trainer = EWCTrainer(ewc_model, lambda_ewc=500.0)
train_vanilla(ewc_model, loader_A, epochs=15)
acc_A_ewc_after_A = evaluate(ewc_model, X_A_test, y_A_test)
ewc_trainer.register_task(loader_A)
ewc_trainer.train_with_ewc(loader_B, epochs=15)
acc_A_ewc_after_B = evaluate(ewc_model, X_A_test, y_A_test)
acc_B_ewc_after_B = evaluate(ewc_model, X_B_test, y_B_test)=== EWC === Task A accuracy after training on A: 0.800 Task A accuracy after training on B (EWC): 0.480 Task B accuracy after training on B (EWC): 0.785 Forgetting on Task A: 0.320
In this particular run, EWC leaves task A forgetting close to the unprotected baseline while still achieving meaningful task B accuracy. The Fisher-weighted penalty changes which parameters the optimizer can move, but the chosen regularization strength does not produce a large retention gain here. This result reflects the method's sensitivity to : some degree of compromise is inevitable when a single parameter vector must serve two different tasks simultaneously.
EWC Importance Distribution
After computing the diagonal Fisher, we can visualize which parameters the model considers most important for task A. High-importance parameters will be strongly protected during task B training.
The distribution reveals an important property: regularization-based methods work because importance is concentrated in a subset of parameters. The isolated bar at contains exact zeros after the numerical floor used for the logarithm; those parameters receive essentially no EWC protection. The continuous mass to its right spans several orders of magnitude, with the right tail representing parameters that most strongly encode task A's solution. If importance were uniformly high, every parameter would require protection and the model could not adapt to new tasks at all. This concentration reflects the fact that neural networks are typically over-parameterized relative to the tasks they solve, leaving many parameters available for repurposing.
SI Implementation
Synaptic Intelligence accumulates importance during training rather than computing it post-hoc. This requires modifying the training loop to track per-parameter gradient-update products.
class SITrainer:
def __init__(self, model, c_si=1.0, xi=0.1):
self.model = model
self.c_si = c_si
self.xi = xi
self.omega = {
n: torch.zeros_like(p) for n, p in model.named_parameters()
}
self.W = {n: torch.zeros_like(p) for n, p in model.named_parameters()}
self.params_prev = {
n: p.data.clone() for n, p in model.named_parameters()
}
def si_loss(self):
loss = 0.0
for n, p in self.model.named_parameters():
delta = self.params_prev[n] - p
loss += (self.omega[n] * delta**2).sum()
return self.c_si * loss
def train_with_si(self, loader, epochs=10, lr=1e-3):
optimizer = optim.SGD(self.model.parameters(), lr=lr)
criterion = nn.BCEWithLogitsLoss()
self.W = {
n: torch.zeros_like(p) for n, p in self.model.named_parameters()
}
params_start = {
n: p.data.clone() for n, p in self.model.named_parameters()
}
self.model.train()
for _ in range(epochs):
for X_batch, y_batch in loader:
params_before = {
n: p.data.clone() for n, p in self.model.named_parameters()
}
optimizer.zero_grad()
task_loss = criterion(self.model(X_batch), y_batch)
reg_loss = self.si_loss()
total_loss = task_loss + reg_loss
total_loss.backward()
for n, p in self.model.named_parameters():
if p.grad is not None:
self.W[n] -= p.grad.data * (p.data - params_before[n])
optimizer.step()
for n, p in self.model.named_parameters():
delta_task = p.data - params_start[n]
self.omega[n] += self.W[n] / (delta_task**2 + self.xi)
self.params_prev = {
n: p.data.clone() for n, p in self.model.named_parameters()
}The most important line in this implementation is self.W[n] -= p.grad.data * (p.data - params_before[n]). This accumulates the negative gradient times the parameter update at each step, approximating the path integral. The subtlety is that we compute the gradient before the optimizer step and the parameter change after, so we must snapshot params_before at the start of each mini-batch. After the full task is done, we normalize self.W by the squared total displacement to get the final importance update, which is added to the running omega accumulation.
si_model = MLP()
si_trainer = SITrainer(si_model, c_si=100.0, xi=0.1)
si_trainer.train_with_si(loader_A, epochs=15, lr=5e-3)
acc_A_si_after_A = evaluate(si_model, X_A_test, y_A_test)
si_trainer.train_with_si(loader_B, epochs=15, lr=5e-3)
acc_A_si_after_B = evaluate(si_model, X_A_test, y_A_test)
acc_B_si_after_B = evaluate(si_model, X_B_test, y_B_test)=== SI === Task A accuracy after training on A: 0.775 Task A accuracy after training on B (SI): 0.545 Task B accuracy after training on B (SI): 0.535 Forgetting on Task A: 0.230
SI also reduces forgetting compared to the baseline. The performance characteristics differ from EWC because the importance estimates reflect the optimization trajectory rather than the final loss curvature. Whether SI or EWC performs better on any particular problem depends on the tasks' structure, the optimizer choice, and the hyperparameter settings.
Comparing Methods
Let's collect the results and compare all three approaches.
method_names = ["No Regularization", "EWC", "SI"]
task_a_after_a = [acc_A_after_A, acc_A_ewc_after_A, acc_A_si_after_A]
task_a_after_b = [acc_A_after_B, acc_A_ewc_after_B, acc_A_si_after_B]
task_b_after_b = [acc_B_after_B, acc_B_ewc_after_B, acc_B_si_after_B]
forgetting_vals = [task_a_after_a[i] - task_a_after_b[i] for i in range(3)]Method Task A->A Task A->B Task B->B Forgetting -------------------------------------------------------------------- No Regularization 0.795 0.470 0.795 0.325 EWC 0.800 0.480 0.785 0.320 SI 0.775 0.545 0.535 0.230
The results are mixed rather than uniformly favorable. SI reduces forgetting clearly in this run, while EWC finishes only slightly below the unprotected baseline. Protecting task A parameters can also limit the model's ability to optimize for task B. This is the fundamental tension in continual learning: plasticity versus stability. The table makes the comparison explicit; read the forgetting and Task B->B columns together to see what each method retains and what it gives up.
The bar charts illustrate the stability-plasticity trade-off in concrete terms. The baseline shows a large drop from the blue bar to the orange bar in the left panel, indicating severe forgetting. EWC's gap is nearly the same at this setting, whereas SI produces a visibly smaller drop. The right panel quantifies that difference: EWC improves the baseline by only 0.005 accuracy points, while SI improves it by 0.095.
Regularization Strength Trade-off
The regularization strength hyperparameter controls the stability-plasticity balance directly. The following plot shows how task A retention and task B performance change as we vary for EWC, using a sweep over different values.
lambda_values = [10.0, 50.0, 200.0, 500.0, 2000.0, 8000.0]
sweep_acc_A = []
sweep_acc_B = []
for lam in lambda_values:
sweep_model = MLP()
train_vanilla(sweep_model, loader_A, epochs=15)
sweep_trainer = EWCTrainer(sweep_model, lambda_ewc=lam)
sweep_trainer.register_task(loader_A)
sweep_trainer.train_with_ewc(loader_B, epochs=15)
sweep_acc_A.append(evaluate(sweep_model, X_A_test, y_A_test))
sweep_acc_B.append(evaluate(sweep_model, X_B_test, y_B_test))
In this sweep, stronger regularization eventually improves task A retention, but the gain is modest until the largest . Task B accuracy is comparatively flat through the middle of the sweep and then drops as the constraint becomes dominant. The curves do not cross, so there is no single visually defined optimum: choosing requires an explicit objective, such as a weighted average of both accuracies or a minimum acceptable score for each task. Similar tasks can tolerate stronger regularization because they benefit from more compatible parameter configurations, while dissimilar tasks may need a lower to preserve enough plasticity for task B.
Key Parameters
The key hyperparameters for both methods are:
- lambda_ewc / c_si: The regularization strength. Higher values preserve more of task A but allow less adaptation to task B. This hyperparameter requires careful tuning; too high and the model cannot learn new tasks, too low and forgetting is not prevented.
- epochs: Training epochs per task. Both methods benefit from sufficient training to reach convergence before computing importance or transitioning to the next task.
- xi (SI only): The numerical stability constant in the SI importance normalization. Small values (0.01 to 0.1) work well in practice.
- learning rate: SI with SGD typically needs a different learning rate than EWC with Adam, since the optimizer choice affects the path integral accumulation in SI.