弾性的重み固定（EWC）（ダンセイテキオモミコテイ）

L(θ) = L_B(θ) + Σᵢ (λ/2) * Fᵢ * (θᵢ - θ*_A,ᵢ)²

import torch
from torch import nn

def compute_fisher_information(model, dataset, device):
    """
    タスクAのデータセットを用いてFisher情報行列を推定
    各パラメータの重要度スコアを返す
    """
    fisher_dict = {n: torch.zeros_like(p)
                   for n, p in model.named_parameters()}
    
    model.eval()
    for inputs, labels in dataset:
        model.zero_grad()
        
        outputs = model(inputs.to(device))
        # 対数尤度の勾配の二乗期待値 = Fisher情報の対角近似
        log_prob = torch.log_softmax(outputs, dim=1)
        loss = log_prob[range(len(labels)), labels].mean()
        loss.backward()
        
        for n, p in model.named_parameters():
            if p.grad is not None:
                fisher_dict[n] += p.grad.data.pow(2)
    
    # データ数で正規化
    for n in fisher_dict:
        fisher_dict[n] /= len(dataset)
    
    return fisher_dict

class EWC:
    def __init__(self, model, dataset, device, importance=1000):
        self.model = model
        self.importance = importance
        
        # タスクA学習後のパラメータを保存
        self.prev_params = {n: p.clone()
                           for n, p in model.named_parameters()}
        
        # Fisher情報を計算
        self.fisher = compute_fisher_information(
            model, dataset, device
        )
    
    def penalty(self):
        """EWCペナルティ項を計算"""
        loss = 0
        for n, p in self.model.named_parameters():
            if n in self.fisher:
                _loss = (self.fisher[n] *
                        (p - self.prev_params[n]).pow(2))
                loss += _loss.sum()
        return self.importance / 2 * loss

# タスクB学習時の使用例
ewc = EWC(model, task_a_dataset, device, importance=1000)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

for batch in task_b_dataloader:
    optimizer.zero_grad()
    
    # タスクBの通常損失
    outputs = model(batch['input_ids'])
    task_loss = criterion(outputs, batch['labels'])
    
    # EWCペナルティを加算
    ewc_loss = ewc.penalty()
    total_loss = task_loss + ewc_loss
    
    total_loss.backward()
    optimizer.step()

# LoRAパラメータのみにEWCを適用
# フルモデルより大幅に計算量削減
lora_params = [(n, p) for n, p in model.named_parameters()
               if 'lora_' in n]

# LoRAアダプタのFisher情報のみ計算
fisher_lora = compute_fisher_for_params(model, dataset, lora_params)

class OnlineEWC:
    def __init__(self, model, gamma=0.95):
        self.model = model
        self.gamma = gamma  # 過去タスクの減衰率
        self.F_accum = {}   # 蓄積Fisher情報
        self.prev_params = {}
    
    def update_after_task(self, dataset, device):
        """タスク完了後にFisher情報を更新"""
        new_fisher = compute_fisher_information(model, dataset, device)
        
        for n in new_fisher:
            if n in self.F_accum:
                # 指数移動平均でFisherを更新
                self.F_accum[n] = (
                    self.gamma * self.F_accum[n] +
                    new_fisher[n]
                )
            else:
                self.F_accum[n] = new_fisher[n]
        
        # 現在のパラメータを保存
        self.prev_params = {
            n: p.clone() for n, p in model.named_parameters()
        }