File size: 4,274 Bytes
f223eaa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """
Bounded Forgetting (R8) — full implementation.
Definitions:
P_j(t) = performance on task j at time t (we use 1/loss).
F_j = max(0, L_after - L_before)
Goal: F_j <= EPS_NEG = 1e-6
Mechanism:
1. Freeze saturated cells (S >= S* for T consecutive steps).
2. Frozen cells: gradient hooks zero out their parameter ranges.
3. New domain → new cells born (disabled in stub), old cells untouched.
"""
import json
from pathlib import Path
from typing import Callable, Dict, List, Optional
import numpy as np
import torch
EPS_NEG = 1e-6
S_STAR = 0.85
T_FREEZE = 3
class CellFreezer:
def __init__(self, S_star: float = S_STAR, T: int = T_FREEZE):
self.S_star = S_star
self.T = T
self.frozen_count = 0
def maybe_freeze(self, cell, current_step: int) -> bool:
if getattr(cell, "frozen", False):
return False
recent = cell.S_history[-self.T:] if hasattr(cell, "S_history") else []
if len(recent) >= self.T and all(s >= self.S_star for s in recent):
cell.frozen = True
cell.freeze_step = current_step
self.frozen_count += 1
return True
return False
def freeze_all_saturated(self, cells, current_step: int) -> int:
n = 0
for c in cells:
if c.state == "active" and self.maybe_freeze(c, current_step):
n += 1
return n
def install_freeze_hooks(model, cells_by_layer):
"""
Register gradient hooks on CellAwareMLP gate/up projections
that zero gradients for frozen cells' neuron ranges.
"""
handles = []
for layer_idx, layer in enumerate(model.model.layers):
mlp = layer.mlp
cells = cells_by_layer.get(layer_idx, [])
if not cells:
continue
def make_hook(cells_ref):
def hook(grad):
g = grad.clone()
for c in cells_ref:
if getattr(c, "frozen", False):
g[..., c.neuron_start:c.neuron_end] = 0.0
return g
return hook
try:
h1 = mlp.gate_proj.weight.register_hook(make_hook(cells))
handles.append(h1)
except AttributeError:
pass
try:
h2 = mlp.up_proj.weight.register_hook(make_hook(cells))
handles.append(h2)
except AttributeError:
pass
return handles
@torch.no_grad()
def evaluate_loss(model, tokens, seq_len=2048, max_batches=10, device="cpu"):
"""Compute average loss on a token stream."""
was_training = model.training
model.eval()
tokens = np.asarray(tokens, dtype=np.int64)
if len(tokens) < seq_len + 1:
if was_training:
model.train()
return float("nan")
n = min(len(tokens) // (seq_len + 1), max_batches)
if n == 0:
if was_training:
model.train()
return float("nan")
total = 0.0
for i in range(n):
s = i * seq_len
x = torch.from_numpy(tokens[s:s + seq_len]).unsqueeze(0).to(device)
y = torch.from_numpy(tokens[s + 1:s + seq_len + 1]).unsqueeze(0).to(device)
out = model(input_ids=x, labels=y)
total += float(out.loss.item())
if was_training:
model.train()
return total / n
def measure_forgetting(model, task_tokens, baseline_losses, device="cpu"):
"""
F_j = max(0, L_after - L_before).
Bounded if all F_j <= EPS_NEG.
"""
results = {}
for task, tokens in task_tokens.items():
l_before = baseline_losses[task]
l_after = evaluate_loss(model, tokens, device=device)
F_j = max(0.0, l_after - l_before)
results[task] = {
"L_before": l_before,
"L_after": l_after,
"F_j": F_j,
"bounded": F_j <= EPS_NEG,
}
return results
def save_forgetting_report(results, path):
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
summary = {
"max_F": max((r["F_j"] for r in results.values()), default=0.0),
"all_bounded": all(r["bounded"] for r in results.values()),
"per_task": results,
}
path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
return path
|