inference_code/inference/pruning/verify.py

"""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