06. Continual Learning and Catastrophic Forgetting — Retraining Without Losing Knowledge

graph TB
    subgraph "Parameter Space ℝⁿ"
        A["θ_A*<br>Task A optimum<br>(92.1% on 10 intents)"]
        B["θ_B*<br>Task B optimum<br>(91.5% on new data but<br>87.3% on old intents)"]
        C["θ_EWC<br>EWC compromise<br>(91.8% old + 89.5% new<br>= 91.1% overall)"]
        D["θ_replay<br>Replay compromise<br>(91.5% old + 90.2% new<br>= 91.2% overall)"]

        A -->|"Naive fine-tune<br>(forgetting path)"| B
        A -->|"EWC λ=1000<br>(constrained path)"| C
        A -->|"Experience Replay<br>(data-balanced path)"| D
    end

    subgraph "Task A Performance Along Path"
        E["Start: 92.1%"]
        F["Naive: drops to 87.3%<br>(-4.8%)"]
        G["EWC: drops to 91.8%<br>(-0.3%)"]
        H["Replay: drops to 91.5%<br>(-0.6%)"]
    end

    style A fill:#c8e6c9
    style B fill:#ffcdd2
    style C fill:#c8e6c9
    style D fill:#c8e6c9

graph TB
    subgraph "Fisher Information Magnitude by Layer (DistilBERT)"
        EMB["Embedding Layer<br>F_avg = 2.3×10⁻⁵<br>→ LOW importance<br>General token representations<br>EWC: allows free movement"]
        L1["Layer 1<br>F_avg = 8.1×10⁻⁵<br>→ LOW importance<br>Basic syntax patterns<br>EWC: mild constraint"]
        L2["Layer 2<br>F_avg = 3.4×10⁻⁴<br>→ MEDIUM importance<br>Phrase-level features<br>EWC: moderate constraint"]
        L3["Layer 3<br>F_avg = 1.2×10⁻³<br>→ MEDIUM-HIGH<br>Semantic grouping<br>EWC: significant constraint"]
        L4["Layer 4<br>F_avg = 5.7×10⁻³<br>→ HIGH importance<br>Intent-discriminative features<br>EWC: strong constraint"]
        L5["Layer 5<br>F_avg = 1.8×10⁻²<br>→ HIGH importance<br>Task-specific representations<br>EWC: strong constraint"]
        L6["Layer 6<br>F_avg = 4.2×10⁻²<br>→ VERY HIGH<br>Classification patterns<br>EWC: very strong constraint"]
        CLS["Classification Head<br>F_avg = 8.5×10⁻²<br>→ CRITICAL<br>Decision boundaries<br>EWC: maximum constraint"]

        EMB --> L1 --> L2 --> L3 --> L4 --> L5 --> L6 --> CLS
    end

    style EMB fill:#e8f5e9
    style L1 fill:#e8f5e9
    style L2 fill:#fff9c4
    style L3 fill:#fff9c4
    style L4 fill:#ffe0b2
    style L5 fill:#ffccbc
    style L6 fill:#ffcdd2
    style CLS fill:#ef9a9a

sequenceDiagram
    participant Batch as New Data Batch<br>"Gift me a manga<br>recommendation"
    participant Model as DistilBERT<br>Parameters θ
    participant CE as CrossEntropy Loss<br>L_B(θ)
    participant EWC as EWC Regularizer<br>λ/2 · Σ Fᵢ(θᵢ-θᵢ*)²
    participant Update as Parameter Update

    Batch->>Model: Forward pass
    Model->>CE: Logits → CE Loss
    Note over CE: ∂L_B/∂θ pushes toward<br>new task optimum<br>(may destroy old knowledge)

    Model->>EWC: Current θ vs stored θ_A*
    Note over EWC: ∂L_EWC/∂θᵢ = λ·Fᵢ·(θᵢ-θᵢ*)<br>Spring force pulling toward<br>old task optimum<br>Strength ∝ Fisher Fᵢ

    CE->>Update: ∂L_B/∂θ (new task gradient)
    EWC->>Update: λ·F·(θ-θ*) (consolidation gradient)

    Note over Update: Total gradient:<br>∂L_B/∂θ + λ·F·(θ-θ*)<br><br>For high-F params: EWC dominates → minimal change<br>For low-F params: CE dominates → free to adapt

    Update->>Model: θ ← θ - η · total_gradient

graph LR
    subgraph "New Task Data (1.2K examples)"
        N1["recommendation: 450"]
        N2["gift_recommendation: 380"]
        N3["product_inquiry: 210"]
        N4["others: 160"]
    end

    subgraph "Replay Buffer (500 from old task)"
        O1["order_status: 55"]
        O2["return: 52"]
        O3["product_inquiry: 50"]
        O4["shipping: 48"]
        O5["complaint: 50"]
        O6["recommendation: 50"]
        O7["account: 48"]
        O8["cancel: 50"]
        O9["exchange: 49"]
        O10["greeting: 48"]
    end

    subgraph "Training Batch (size 32)"
        B1["20 from new task (62.5%)"]
        B2["12 from replay buffer (37.5%)"]
    end

    N1 --> B1
    N2 --> B1
    O1 --> B2
    O2 --> B2
    O3 --> B2

    B1 --> LOSS["Combined CE Loss<br>= L_new + μ·L_replay<br>μ = 1.0 (equal weight)"]
    B2 --> LOSS

    style B1 fill:#fff9c4
    style B2 fill:#bbdefb

graph TB
    subgraph "Month 1: Initial Training"
        M1["Train on 5K labeled examples<br>10 intents, 92.1% accuracy<br>Save θ₁* and F₁"]
    end

    subgraph "Month 2: Seasonal Update"
        M2A["New data: 1.2K (gift reco spike)"]
        M2B["Replay buffer: 500 (from Month 1)"]
        M2C["EWC: λ₁·F₁·(θ-θ₁*)²"]
        M2D["Result: 91.2% overall + gift_reco<br>Save θ₂* and F₂ = 0.9·F₁ + F₂_new"]
        M2A --> M2D
        M2B --> M2D
        M2C --> M2D
    end

    subgraph "Month 3: Trend Shift"
        M3A["New data: 800 (anime adaptation queries)"]
        M3B["Replay buffer: 500 (from Months 1+2)"]
        M3C["Online EWC: λ·F₂·(θ-θ₂*)²"]
        M3D["Result: 90.8% on old + new patterns<br>Save θ₃* and F₃ = 0.9·F₂ + F₃_new"]
        M3A --> M3D
        M3B --> M3D
        M3C --> M3D
    end

    M1 --> M2A
    M2D --> M3A

    style M1 fill:#c8e6c9
    style M2D fill:#c8e6c9
    style M3D fill:#c8e6c9

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import copy
from collections import defaultdict


class EWCTrainer:
    """
    Elastic Weight Consolidation for continual learning.

    Computes Fisher Information from old-task data, then adds a
    quadratic penalty when training on new data.
    """

    def __init__(
        self,
        model: nn.Module,
        fisher_sample_size: int = 500,
        ewc_lambda: float = 1000.0,
        online_gamma: float = 0.9,
    ):
        self.model = model
        self.fisher_sample_size = fisher_sample_size
        self.ewc_lambda = ewc_lambda
        self.online_gamma = online_gamma

        # Store consolidated Fisher and optimal params
        self.fisher: dict[str, torch.Tensor] = {}
        self.optimal_params: dict[str, torch.Tensor] = {}
        self.task_count = 0

    def compute_fisher(self, dataloader: DataLoader):
        """
        Compute empirical Fisher Information Matrix (diagonal approximation).

        F_i = E[(∂ log p(y|x,θ) / ∂θ_i)²]

        Uses squared gradients from negative log-likelihood over old data samples.
        """
        self.model.eval()
        fisher = {
            name: torch.zeros_like(param)
            for name, param in self.model.named_parameters()
            if param.requires_grad
        }

        n_samples = 0
        for batch in dataloader:
            if n_samples >= self.fisher_sample_size:
                break

            input_ids = batch["input_ids"].to(self.model.device)
            attention_mask = batch["attention_mask"].to(self.model.device)
            labels = batch["labels"].to(self.model.device)

            self.model.zero_grad()
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=labels,
            )
            loss = outputs.loss
            loss.backward()

            for name, param in self.model.named_parameters():
                if param.requires_grad and param.grad is not None:
                    fisher[name] += param.grad.data ** 2

            n_samples += input_ids.shape[0]

        # Average over samples
        for name in fisher:
            fisher[name] /= n_samples

        return fisher

    def consolidate(self, dataloader: DataLoader):
        """
        After training on a task, consolidate knowledge:
        1. Compute Fisher Information on current task
        2. Store optimal parameters
        3. If online EWC, merge with previous Fisher
        """
        new_fisher = self.compute_fisher(dataloader)

        if self.task_count == 0:
            # First task — just store
            self.fisher = new_fisher
        else:
            # Online EWC: merge with running estimate
            for name in new_fisher:
                self.fisher[name] = (
                    self.online_gamma * self.fisher[name] + new_fisher[name]
                )

        # Store current optimal parameters
        self.optimal_params = {
            name: param.data.clone()
            for name, param in self.model.named_parameters()
            if param.requires_grad
        }

        self.task_count += 1

    def ewc_loss(self) -> torch.Tensor:
        """
        Compute EWC penalty: (λ/2) Σ_i F_i (θ_i - θ_i*)²

        Returns 0 if no task has been consolidated yet.
        """
        if self.task_count == 0:
            return torch.tensor(0.0, device=self.model.device)

        loss = torch.tensor(0.0, device=self.model.device)
        for name, param in self.model.named_parameters():
            if name in self.fisher:
                loss += (
                    self.fisher[name] * (param - self.optimal_params[name]) ** 2
                ).sum()

        return (self.ewc_lambda / 2) * loss

    def train_step(self, batch, optimizer):
        """Single training step with EWC regularization."""
        self.model.train()

        input_ids = batch["input_ids"].to(self.model.device)
        attention_mask = batch["attention_mask"].to(self.model.device)
        labels = batch["labels"].to(self.model.device)

        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
        )

        task_loss = outputs.loss
        ewc_penalty = self.ewc_loss()
        total_loss = task_loss + ewc_penalty

        optimizer.zero_grad()
        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
        optimizer.step()

        return {
            "total_loss": total_loss.item(),
            "task_loss": task_loss.item(),
            "ewc_penalty": ewc_penalty.item(),
        }

import random
from collections import defaultdict


class ReplayBuffer:
    """
    Stratified replay buffer that maintains balanced class representation.
    Uses reservoir sampling for efficient memory-bounded updates.
    """

    def __init__(self, max_per_class: int = 50):
        self.max_per_class = max_per_class
        self.buffer: dict[int, list] = defaultdict(list)
        self.counts: dict[int, int] = defaultdict(int)

    def add_examples(self, examples: list[dict]):
        """
        Add examples using reservoir sampling.
        Ensures each class has at most max_per_class examples.
        """
        for example in examples:
            label = example["label"]
            self.counts[label] += 1

            if len(self.buffer[label]) < self.max_per_class:
                self.buffer[label].append(example)
            else:
                # Reservoir sampling: replace with probability max_per_class / count
                idx = random.randint(0, self.counts[label] - 1)
                if idx < self.max_per_class:
                    self.buffer[label][idx] = example

    def sample(self, n: int) -> list[dict]:
        """Sample n examples with uniform class distribution."""
        all_classes = list(self.buffer.keys())
        per_class = max(1, n // len(all_classes))
        samples = []
        for cls in all_classes:
            cls_samples = random.sample(
                self.buffer[cls],
                min(per_class, len(self.buffer[cls])),
            )
            samples.extend(cls_samples)
        random.shuffle(samples)
        return samples[:n]

    @property
    def total_size(self) -> int:
        return sum(len(v) for v in self.buffer.values())

    def stats(self) -> dict:
        return {cls: len(examples) for cls, examples in self.buffer.items()}


class ContinualTrainer:
    """
    Combines EWC regularization with experience replay.
    This hybrid approach outperforms either technique alone.
    """

    def __init__(
        self,
        model,
        ewc_lambda: float = 500.0,
        replay_ratio: float = 0.3,
        replay_buffer_per_class: int = 50,
    ):
        self.model = model
        self.ewc = EWCTrainer(model, ewc_lambda=ewc_lambda)
        self.replay_buffer = ReplayBuffer(max_per_class=replay_buffer_per_class)
        self.replay_ratio = replay_ratio

    def train_on_new_task(
        self,
        train_data,
        epochs: int = 5,
        lr: float = 2e-5,
    ):
        """
        Train on new task data with EWC + replay.

        Each batch is constructed as:
        - (1 - replay_ratio) from new data
        - replay_ratio from replay buffer
        """
        optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr)

        new_batch_size = int(32 * (1 - self.replay_ratio))  # 22
        replay_batch_size = 32 - new_batch_size              # 10

        for epoch in range(epochs):
            epoch_loss = {"total": 0, "task": 0, "ewc": 0}
            n_steps = 0

            new_loader = DataLoader(train_data, batch_size=new_batch_size, shuffle=True)

            for new_batch in new_loader:
                # Construct mixed batch
                if self.replay_buffer.total_size > 0:
                    replay_samples = self.replay_buffer.sample(replay_batch_size)
                    batch = self._merge_batches(new_batch, replay_samples)
                else:
                    batch = new_batch

                metrics = self.ewc.train_step(batch, optimizer)
                epoch_loss["total"] += metrics["total_loss"]
                epoch_loss["task"] += metrics["task_loss"]
                epoch_loss["ewc"] += metrics["ewc_penalty"]
                n_steps += 1

            for k in epoch_loss:
                epoch_loss[k] /= max(n_steps, 1)

            print(
                f"Epoch {epoch+1}: total={epoch_loss['total']:.4f} "
                f"task={epoch_loss['task']:.4f} ewc={epoch_loss['ewc']:.4f}"
            )

        # After training, consolidate and update replay buffer
        self.ewc.consolidate(DataLoader(train_data, batch_size=32))
        self.replay_buffer.add_examples(
            [{"input_ids": x["input_ids"], "attention_mask": x["attention_mask"],
              "label": x["labels"]} for x in train_data]
        )

    def _merge_batches(self, new_batch, replay_samples):
        """Merge new data batch with replay buffer samples."""
        replay_ids = torch.stack([s["input_ids"] for s in replay_samples])
        replay_masks = torch.stack([s["attention_mask"] for s in replay_samples])
        replay_labels = torch.tensor([s["label"] for s in replay_samples])

        return {
            "input_ids": torch.cat([new_batch["input_ids"], replay_ids]),
            "attention_mask": torch.cat([new_batch["attention_mask"], replay_masks]),
            "labels": torch.cat([new_batch["labels"], replay_labels]),
        }

import json
from datetime import datetime
import mlflow


def monthly_retrain_pipeline(
    model_path: str,
    new_data_path: str,
    old_data_path: str,
    fisher_path: str = None,
):
    """
    MangaAssist monthly retraining pipeline with catastrophic forgetting prevention.

    Steps:
    1. Load current production model and Fisher (if exists)
    2. Detect if new intents have emerged
    3. Train with EWC + replay
    4. Validate against old-task benchmarks
    5. Deploy only if old-task regression < 1%
    """
    model = AutoModelForSequenceClassification.from_pretrained(model_path)
    tokenizer = AutoTokenizer.from_pretrained(model_path)

    # Initialize continual trainer
    trainer = ContinualTrainer(
        model=model,
        ewc_lambda=500.0,
        replay_ratio=0.3,
        replay_buffer_per_class=50,
    )

    # Load previous Fisher if available
    if fisher_path:
        fisher_data = torch.load(fisher_path)
        trainer.ewc.fisher = fisher_data["fisher"]
        trainer.ewc.optimal_params = fisher_data["optimal_params"]
        trainer.ewc.task_count = fisher_data["task_count"]

    # Load replay buffer from previous tasks
    if old_data_path:
        old_data = load_dataset(old_data_path)
        trainer.replay_buffer.add_examples(old_data)

    # Train on new data
    new_data = load_dataset(new_data_path)
    trainer.train_on_new_task(new_data, epochs=5, lr=2e-5)

    # Validation gate: check old-task regression
    old_eval = evaluate_model(model, tokenizer, old_data_path)
    new_eval = evaluate_model(model, tokenizer, new_data_path)

    with mlflow.start_run():
        mlflow.log_metric("old_task_accuracy", old_eval["accuracy"])
        mlflow.log_metric("new_task_accuracy", new_eval["accuracy"])
        mlflow.log_metric("combined_accuracy",
                         0.8 * old_eval["accuracy"] + 0.2 * new_eval["accuracy"])

    # Deployment gate
    REGRESSION_THRESHOLD = 0.01  # Max 1% old-task accuracy drop
    if old_eval["accuracy"] >= (0.921 - REGRESSION_THRESHOLD):
        print(f"✓ Old-task accuracy {old_eval['accuracy']:.4f} within threshold")
        model.save_pretrained(f"./models/intent_v{datetime.now():%Y%m}")
        # Save Fisher for next cycle
        torch.save({
            "fisher": trainer.ewc.fisher,
            "optimal_params": trainer.ewc.optimal_params,
            "task_count": trainer.ewc.task_count,
        }, f"./models/fisher_v{datetime.now():%Y%m}.pt")
    else:
        raise ValueError(
            f"✗ Old-task accuracy {old_eval['accuracy']:.4f} below threshold "
            f"{0.921 - REGRESSION_THRESHOLD:.4f}. Rejecting model update."
        )