Phase 33.1 — ElasticWeightConsolidator: Fisher Information-based parameter protection for continual learning

[web3guru888]

Overview

The ElasticWeightConsolidator implements Elastic Weight Consolidation (EWC) and related regularization-based continual learning methods. It computes Fisher Information matrices to identify task-critical parameters and applies quadratic penalties during subsequent task learning, preventing catastrophic forgetting while allowing plasticity on less important weights.

Key Responsibilities

• Fisher Information computation — Diagonal and block-diagonal Fisher approximation from task-specific data
• Parameter importance mapping — Online importance estimation using both EWC (Kirkpatrick et al., 2017) and Synaptic Intelligence (Zenke et al., 2017)
• Task-specific regularization — Quadratic penalty terms anchoring important parameters to their post-training values
• Online Fisher updates — Incremental Fisher matrix accumulation across tasks without storing full task datasets
• Memory Aware Synapses (MAS) — Unsupervised importance estimation via gradient magnitude (Aljundi et al., 2018)
• Multi-task consolidation — Merging importance maps across task sequence with configurable decay
• Importance visualization — Heatmaps of parameter importance across layers and tasks

Interfaces

Inputs

• model: nn.Module — Neural network whose parameters are being consolidated
• task_data: DataLoader — Task-specific data for Fisher computation
• task_id: str — Unique task identifier
• lambda_ewc: float — Regularization strength (default: 5000.0)
• fisher_method: Literal["diagonal", "block_diagonal", "kfac"] — Fisher approximation type
• online: bool — Whether to use online (running) Fisher estimation

Outputs

• ConsolidationResult — Contains penalty loss, per-layer importance stats, forgetting risk score
• ImportanceMap — Dict[str, Tensor] mapping parameter names to importance scores
• FisherSnapshot — Serializable Fisher matrix for checkpointing

Acceptance Criteria

• Fisher Information matrix computed within 2x forward-pass time
• EWC penalty prevents >5% accuracy degradation on previous tasks (Split-MNIST benchmark)
• Online Fisher updates require O(1) additional memory per task
• MAS importance estimation works without task labels
• Supports models with >100M parameters efficiently
• Importance maps serializable and loadable across sessions
• Integration tests with Phase 33.5 ContinualOrchestrator
• 95%+ test coverage, all public methods typed

References

• Kirkpatrick et al. (2017) — Overcoming Catastrophic Forgetting (EWC)
• Zenke et al. (2017) — Synaptic Intelligence (SI)
• Aljundi et al. (2018) — Memory Aware Synapses (MAS)
• Schwarz et al. (2018) — Progress & Compress
• Planning: Phase 33 Planning — Continual Learning & Catastrophic Forgetting Prevention #690

[web3guru888]

Implementation Notes — ElasticWeightConsolidator

Python Skeleton

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Dict, Iterator, Literal, Optional, Tuple
import torch
from torch import Tensor, nn
from torch.utils.data import DataLoader

@dataclass(frozen=True)
class ConsolidationResult:
    """Output of a consolidation step."""
    penalty_loss: float
    per_layer_importance: Dict[str, float]
    forgetting_risk: float          # 0.0 = safe, 1.0 = high risk

@dataclass
class FisherSnapshot:
    """Serialisable Fisher information for one task."""
    task_id: str
    diagonal: Dict[str, Tensor]     # param_name → diagonal Fisher
    optimal_params: Dict[str, Tensor]
    n_samples: int

class ElasticWeightConsolidator:
    """EWC / SI / MAS continual-learning regulariser."""

    def __init__(
        self,
        model: nn.Module,
        lambda_ewc: float = 5000.0,
        fisher_method: Literal["diagonal", "block_diagonal", "kfac"] = "diagonal",
        online: bool = False,
        gamma: float = 0.95,          # online Fisher decay
    ) -> None:
        self._model = model
        self._lambda = lambda_ewc
        self._method = fisher_method
        self._online = online
        self._gamma = gamma
        self._snapshots: list[FisherSnapshot] = []
        self._running_fisher: Dict[str, Tensor] = {}

    # ── public API ──────────────────────────────────────────────
    def compute_fisher(
        self, task_id: str, loader: DataLoader, n_samples: int = 2000
    ) -> FisherSnapshot:
        """Estimate diagonal Fisher information from *loader*."""
        fisher: Dict[str, Tensor] = {
            n: torch.zeros_like(p) for n, p in self._named_params()
        }
        self._model.eval()
        count = 0
        for x, y in loader:
            if count >= n_samples:
                break
            log_probs = self._model(x).log_softmax(dim=-1)
            labels = log_probs.argmax(dim=-1)            # empirical Fisher
            loss = nn.functional.nll_loss(log_probs, labels)
            loss.backward()
            for n, p in self._named_params():
                if p.grad is not None:
                    fisher[n] += p.grad.detach() ** 2
            self._model.zero_grad()
            count += x.size(0)
        # normalise
        for n in fisher:
            fisher[n] /= count
        snap = FisherSnapshot(
            task_id=task_id,
            diagonal=fisher,
            optimal_params={n: p.detach().clone() for n, p in self._named_params()},
            n_samples=count,
        )
        if self._online:
            self._update_running_fisher(fisher)
        self._snapshots.append(snap)
        return snap

    def penalty(self) -> ConsolidationResult:
        """Compute EWC quadratic penalty over all stored snapshots."""
        total = torch.tensor(0.0)
        layer_imp: Dict[str, float] = {}
        for snap in self._snapshots:
            for n, p in self._named_params():
                diff = (p - snap.optimal_params[n]) ** 2
                contribution = (snap.diagonal[n] * diff).sum()
                total = total + contribution
                layer_imp[n] = layer_imp.get(n, 0.0) + contribution.item()
        penalty_loss = self._lambda * total
        risk = min(1.0, penalty_loss.item() / (self._lambda * 10))
        return ConsolidationResult(
            penalty_loss=penalty_loss.item(),
            per_layer_importance=layer_imp,
            forgetting_risk=risk,
        )

    def mas_importance(self, loader: DataLoader) -> Dict[str, Tensor]:
        """Memory-Aware-Synapses unsupervised importance (Aljundi 2018)."""
        importance: Dict[str, Tensor] = {
            n: torch.zeros_like(p) for n, p in self._named_params()
        }
        self._model.eval()
        for x, _ in loader:
            out = self._model(x)
            l2 = out.norm(dim=-1).mean()
            l2.backward()
            for n, p in self._named_params():
                if p.grad is not None:
                    importance[n] += p.grad.detach().abs()
            self._model.zero_grad()
        return importance

    # ── internals ───────────────────────────────────────────────
    def _named_params(self) -> Iterator[Tuple[str, nn.Parameter]]:
        return ((n, p) for n, p in self._model.named_parameters() if p.requires_grad)

    def _update_running_fisher(self, new: Dict[str, Tensor]) -> None:
        for n, f in new.items():
            if n in self._running_fisher:
                self._running_fisher[n] = self._gamma * self._running_fisher[n] + f
            else:
                self._running_fisher[n] = f.clone()

Algorithm — Online EWC (Schwarz et al. 2018)

procedure ONLINE_EWC_UPDATE(model, task_data, γ):
    F_new ← compute_diagonal_fisher(model, task_data)
    θ*    ← current model parameters
    for each parameter p:
        F_running[p] ← γ · F_running[p] + F_new[p]    // exponential moving average
        θ*_stored[p] ← θ*[p]                            // anchor point
    return F_running, θ*_stored

procedure EWC_LOSS(model, F_running, θ*_stored, λ):
    penalty ← 0
    for each parameter p:
        penalty += F_running[p] · (model[p] - θ*_stored[p])²
    return λ · penalty

Type Table (mypy-compatible)

Symbol	Type	Description
model	nn.Module	Target model
lambda_ewc	float	Regularization strength
fisher_method	Literal["diagonal", "block_diagonal", "kfac"]	Approximation type
online	bool	Online vs. offline EWC
gamma	float	Online Fisher decay rate
ConsolidationResult.penalty_loss	float	Scalar penalty value
ConsolidationResult.per_layer_importance	Dict[str, float]	Per-layer importance
ConsolidationResult.forgetting_risk	float	Risk score [0, 1]
FisherSnapshot.diagonal	Dict[str, Tensor]	Diagonal Fisher entries
FisherSnapshot.optimal_params	Dict[str, Tensor]	Anchor parameters