| """ |
| ewc.py — Elastic Weight Consolidation for KAIZEN forgetting benchmark. |
| |
| Kirkpatrick et al. 2017: "Overcoming catastrophic forgetting in neural networks." |
| L_total = L_task + λ * Σ_t Σ_i F_{t,i} * (θ_i - θ*_{t,i})² |
| |
| Applied to LoRAAdapter parameters only. Base model is frozen throughout. |
| Fisher Information: empirical Fisher (mean squared gradient) computed on |
| domain j training data immediately after convergence. |
| """ |
| import torch |
|
|
| from lora import LoRAAdapter, KaizenWithLoRA |
| from online_learner import build_update_seq |
|
|
|
|
| class EWC: |
| """Diagonal Laplace EWC regularizer for LoRAAdapter. |
| |
| Usage: |
| ewc = EWC(lambda_ewc=1000.0) |
| # After training domain j: |
| ewc.update(model, adapter, domain_j_tasks, tokenizer) |
| # During domain j+1 training, add ewc.penalty(adapter) to task loss. |
| """ |
|
|
| def __init__(self, lambda_ewc: float = 1000.0): |
| self.lambda_ewc = lambda_ewc |
| self._fishers: list[dict[str, torch.Tensor]] = [] |
| self._optimal: list[dict[str, torch.Tensor]] = [] |
|
|
| def update(self, model: KaizenWithLoRA, adapter: LoRAAdapter, |
| tasks: list, tokenizer) -> float: |
| """Compute empirical Fisher + save parameter anchor for domain tasks. |
| |
| Call AFTER adapter has converged on domain j, BEFORE training domain j+1. |
| Returns mean Fisher L2 norm (diagnostic — useful for tuning lambda_ewc). |
| """ |
| fisher = {n: torch.zeros_like(p) for n, p in adapter.named_parameters()} |
| n_seen = 0 |
|
|
| for p in model.parameters(): |
| p.requires_grad_(False) |
| for p in adapter.parameters(): |
| p.requires_grad_(True) |
|
|
| for question, answer, _ in tasks: |
| x, y = build_update_seq(tokenizer, question, answer) |
| adapter.zero_grad() |
| _, loss = model(x, targets=y, adapter=adapter) |
| loss.backward() |
| for n, p in adapter.named_parameters(): |
| if p.grad is not None: |
| fisher[n] = fisher[n] + p.grad.data.pow(2) |
| n_seen += 1 |
|
|
| scale = 1.0 / max(n_seen, 1) |
| anchored_fisher = {n: (f * scale).clone() for n, f in fisher.items()} |
| self._fishers.append(anchored_fisher) |
| self._optimal.append({n: p.data.clone() for n, p in adapter.named_parameters()}) |
|
|
| mean_fisher_norm = float( |
| torch.stack([f.norm() for f in anchored_fisher.values()]).mean().item() |
| ) |
| return mean_fisher_norm |
|
|
| def penalty(self, adapter: LoRAAdapter) -> torch.Tensor: |
| """EWC regularization term: λ * Σ_t Σ_i F_{t,i} * (θ_i - θ*_{t,i})² |
| |
| Differentiable w.r.t. adapter.parameters(). Returns scalar tensor. |
| Zero (no grad cost) when no anchor stored yet (first domain). |
| """ |
| if not self._fishers: |
| return torch.tensor(0.0) |
|
|
| penalties = [] |
| for fisher_t, opt_t in zip(self._fishers, self._optimal): |
| for n, p in adapter.named_parameters(): |
| penalties.append((fisher_t[n] * (p - opt_t[n]).pow(2)).sum()) |
|
|
| return self.lambda_ewc * torch.stack(penalties).sum() |
|
|
