File size: 8,866 Bytes
fc3971e | 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 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 | """Offline diagnostics for signed TENP vs actual removal damage."""
from __future__ import annotations
import torch
from .collector import perturbation_sensitivity, stacked_token_cache
from .score import PROJ_EPS
def spearman_corr(x: torch.Tensor, y: torch.Tensor, eps: float = 1e-8) -> float:
"""Average-rank Spearman. Fine for distinct scores; ties share mid-ranks."""
a = x.detach().float().reshape(-1)
b = y.detach().float().reshape(-1)
if a.numel() != b.numel() or a.numel() < 3:
return float("nan")
ra = _average_rank(a)
rb = _average_rank(b)
ra = ra - ra.mean()
rb = rb - rb.mean()
denom = (ra.norm() * rb.norm()).clamp_min(eps)
return float((ra * rb).sum() / denom)
def _average_rank(x: torch.Tensor) -> torch.Tensor:
n = int(x.numel())
order = torch.argsort(x)
ranks = torch.empty(n, device=x.device, dtype=torch.float32)
sorted_x = x[order]
i = 0
while i < n:
j = i
while j + 1 < n and float(sorted_x[j + 1]) == float(sorted_x[i]):
j += 1
mid = 0.5 * (i + j) + 1.0
ranks[order[i : j + 1]] = mid
i = j + 1
return ranks
def single_neuron_removal_damage(
rel_abs_z: torch.Tensor,
n_tokens: int,
down_weight: torch.Tensor,
) -> torch.Tensor:
"""Mean relative L2 of masking neuron i at expert output (top-1 MoE-equivalent).
During collect we store sum_t |z_{t,i}| / ||y_t||. Then:
Δ_i = (rel_abs_z_i / N) * ||W_down[:, i]||_2
which equals E_t ||z_i W_i|| / ||y||.
"""
if n_tokens <= 0:
return torch.zeros_like(rel_abs_z, dtype=torch.float32)
col_norm = down_weight.detach().float().norm(dim=0)
return rel_abs_z.detach().float() / float(n_tokens) * col_norm
def group_mean_damage(damage: torch.Tensor, idx: torch.Tensor) -> float:
if idx.numel() == 0:
return float("nan")
return float(damage[idx.long()].mean())
def group_removal_relative_l2(
z: torch.Tensor,
y: torch.Tensor,
down_weight: torch.Tensor,
idx: torch.Tensor,
eps: float = PROJ_EPS,
) -> float:
"""Exact mean ||sum_{i in G} z_i W_i|| / ||y|| on cached tokens."""
if idx is None or int(idx.numel()) == 0 or z.numel() == 0:
return float("nan")
keep = idx.detach().long().to(device=z.device)
w = down_weight.detach().to(device=z.device, dtype=torch.float32)
z_f = z.float()
y_n = y.float().norm(dim=-1).clamp_min(eps)
delta = z_f.index_select(1, keep) @ w.index_select(1, keep).T
return float((delta.norm(dim=-1) / y_n).mean())
def sensitivity_by_bin(
cache_z: list,
cache_y: list,
down_weight: torch.Tensor,
ratios: tuple[float, ...] = (0.05, 0.10),
repeats: int = 3,
seed: int = 0,
) -> dict[str, list[float]]:
"""Offline D^ρ per timestep bin from the token reservoir."""
gen = torch.Generator(device="cpu")
gen.manual_seed(int(seed))
out: dict[str, list[float]] = {f"{rho:.2f}": [] for rho in ratios}
n_bins = max(len(cache_z or []), len(cache_y or []))
for bin_id in range(n_bins):
z = None if cache_z is None or bin_id >= len(cache_z) else cache_z[bin_id]
y = None if cache_y is None or bin_id >= len(cache_y) else cache_y[bin_id]
for rho in ratios:
key = f"{rho:.2f}"
if z is None or y is None or int(z.numel()) == 0:
out[key].append(float("nan"))
continue
out[key].append(
perturbation_sensitivity(z, y, down_weight, rho, repeats, gen)
)
return out
def redundancy_pair_causal(
z: torch.Tensor,
y: torch.Tensor,
down_weight: torch.Tensor,
score: torch.Tensor,
r_plus: torch.Tensor | None = None,
gram: torch.Tensor | None = None,
n_cand: int = 256,
n_pairs: int = 16,
seed: int = 0,
) -> dict:
"""Compare high-R vs low-R pairs: additivity = Δij / (Δi+Δj).
Same-direction (redundant) pairs add constructively, so additivity is
closer to 1 than for near-orthogonal controls with similar TENP scores.
"""
from .redundancy import contribution_redundancy
s = score.detach().float().reshape(-1)
pos = (s > 0).nonzero(as_tuple=False).reshape(-1)
if int(pos.numel()) < 4 or z.numel() == 0:
return {"skipped": "need cached tokens and >=4 positive-score neurons"}
k = min(int(n_cand), int(pos.numel()))
_, local = torch.topk(s[pos], k)
idx = pos[local]
if r_plus is not None:
sub = r_plus.detach().float()[idx][:, idx].clone()
else:
g_sub = None if gram is None else gram.detach().float()[idx][:, idx]
sub = contribution_redundancy(g_sub, down_weight.detach().float()[:, idx])
sub.fill_diagonal_(0)
tri_i, tri_j = torch.triu_indices(k, k, offset=1)
r_flat = sub[tri_i, tri_j]
if int(r_flat.numel()) < 2:
return {"skipped": "not enough pairs"}
n_take = min(int(n_pairs), int(r_flat.numel()))
high_sel = torch.topk(r_flat, n_take).indices
low_sel = torch.topk(-r_flat, n_take).indices
def _pack(sel: torch.Tensor) -> dict:
additivities = []
overlaps = []
r_vals = []
for t in sel.tolist():
li = int(tri_i[t])
lj = int(tri_j[t])
a = int(idx[li].item())
b = int(idx[lj].item())
di = group_removal_relative_l2(z, y, down_weight, torch.tensor([a]))
dj = group_removal_relative_l2(z, y, down_weight, torch.tensor([b]))
both = group_removal_relative_l2(z, y, down_weight, torch.tensor([a, b]))
denom = di + dj
additivity = float("nan") if denom <= 0 else float(both / denom)
overlap = float("nan") if denom <= 0 else float((denom - both) / denom)
additivities.append(additivity)
overlaps.append(overlap)
r_vals.append(float(sub[li, lj]))
finite_a = [x for x in additivities if x == x]
finite_o = [x for x in overlaps if x == x]
return {
"n_pairs": len(additivities),
"r_mean": round(float(sum(r_vals) / max(len(r_vals), 1)), 4),
"additivity_mean": (
round(float(sum(finite_a) / len(finite_a)), 4) if finite_a else float("nan")
),
"overlap_mean": (
round(float(sum(finite_o) / len(finite_o)), 4) if finite_o else float("nan")
),
}
del seed
high = _pack(high_sel)
low = _pack(low_sel)
return {
"high_r": high,
"low_r": low,
"additivity_gap": _finite_gap(high, low, "additivity_mean"),
"overlap_gap": _finite_gap(high, low, "overlap_mean"),
"note": "aligned/redundant pairs have additivity Δij/(Δi+Δj) closer to 1",
}
def _finite_gap(high: dict, low: dict, key: str) -> float:
a = high.get(key)
b = low.get(key)
if not isinstance(a, float) or not isinstance(b, float):
return float("nan")
if a != a or b != b:
return float("nan")
return round(float(a - b), 4)
def cache_group_report(
stats,
down_weight: torch.Tensor,
signed: torch.Tensor,
damage: torch.Tensor,
frac: float = 0.10,
) -> dict:
"""Exact group-removal Δ on the reservoir vs mean of single-neuron Δ."""
n_bot = max(1, int(round(float(frac) * int(signed.numel()))))
bot = torch.topk(signed, n_bot, largest=False).indices
top = torch.topk(signed, n_bot, largest=True).indices
packed = stacked_token_cache(stats)
rec = {
"approx_mean_delta": {
"lowest_signed_10pct": round(group_mean_damage(damage, bot), 6),
"highest_signed_10pct": round(group_mean_damage(damage, top), 6),
}
}
if packed is None:
rec["exact_group_delta"] = {"skipped": "missing token cache; re-collect"}
rec["sensitivity_rho"] = {"skipped": "missing token cache; re-collect"}
rec["redundancy_pairs"] = {"skipped": "missing token cache; re-collect"}
return rec
z, y, _gate = packed
rec["n_cached_tokens"] = int(z.shape[0])
rec["exact_group_delta"] = {
"lowest_signed_10pct": round(
group_removal_relative_l2(z, y, down_weight, bot), 6
),
"highest_signed_10pct": round(
group_removal_relative_l2(z, y, down_weight, top), 6
),
}
rec["sensitivity_rho"] = sensitivity_by_bin(
stats.cache_z or [],
stats.cache_y or [],
down_weight,
ratios=(0.05, 0.10),
)
d5 = torch.tensor(rec["sensitivity_rho"]["0.05"], dtype=torch.float32)
d10 = torch.tensor(rec["sensitivity_rho"]["0.10"], dtype=torch.float32)
live = torch.isfinite(d5) & torch.isfinite(d10)
rec["spearman_d05_vs_d10"] = (
round(spearman_corr(d5[live], d10[live]), 4) if int(live.sum()) >= 3 else float("nan")
)
return rec
|