| """Contribution-space redundancy and diversity-aware keep selection. |
| |
| Negative TENP scores are never diversity-discounted: a neuron that points |
| against the expert output stays a prune candidate even if it looks unique. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
|
|
| from .surgery import round_to_multiple |
|
|
|
|
| def n_keep_width(d_int: int, p_mid: float, multiple: int = 4) -> int: |
| n_keep = max(multiple, round_to_multiple(int((1.0 - p_mid) * d_int), multiple)) |
| return min(n_keep, d_int) |
|
|
|
|
| def gram_cosine(gram: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: |
| """Column cosine from a Gram matrix G ≈ AᵀA, shape [d, d].""" |
| if gram.ndim != 2 or gram.shape[0] != gram.shape[1]: |
| raise ValueError(f"gram must be square, got {tuple(gram.shape)}") |
| g = gram.detach().float() |
| g = (g + g.T) * 0.5 |
| nrm = torch.diag(g).clamp_min(0).sqrt().clamp_min(eps) |
| return g / (nrm.unsqueeze(0) * nrm.unsqueeze(1)) |
|
|
|
|
| def weight_cosine(down_weight: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: |
| """Column cosine of W_down, shape [hidden, d_int] -> [d_int, d_int].""" |
| w = down_weight.detach().float() |
| nrm = w.norm(dim=0).clamp_min(eps) |
| hat = w / nrm.unsqueeze(0) |
| return hat.T @ hat |
|
|
|
|
| def contribution_redundancy( |
| gram: torch.Tensor | None, |
| down_weight: torch.Tensor, |
| eps: float = 1e-8, |
| ) -> torch.Tensor: |
| """R⁺ = max(0, cos(a) cos(W)). Missing Gram falls back to weight cosine only.""" |
| r_w = weight_cosine(down_weight, eps=eps) |
| if gram is None: |
| r = r_w |
| else: |
| r = gram_cosine(gram, eps=eps) * r_w |
| r = torch.maximum(r, torch.zeros_like(r)) |
| d = int(r.shape[0]) |
| eye = torch.eye(d, device=r.device, dtype=r.dtype) |
| return r * (1.0 - eye) + eye |
|
|
|
|
| def greedy_diversity_keep( |
| score: torch.Tensor, |
| r_plus: torch.Tensor, |
| n_keep: int, |
| ) -> torch.Tensor: |
| """Greedy keep: U_i = S_i (1 - r_i^max) for S_i > 0, else U_i = S_i.""" |
| s = score.detach().float().reshape(-1) |
| d = int(s.numel()) |
| n_keep = min(max(int(n_keep), 1), d) |
| if r_plus.shape != (d, d): |
| raise ValueError(f"r_plus must be [{d}, {d}], got {tuple(r_plus.shape)}") |
| r = r_plus.detach().float() |
| remaining = torch.ones(d, dtype=torch.bool, device=s.device) |
| r_max = torch.zeros(d, device=s.device, dtype=s.dtype) |
| chosen: list[int] = [] |
| neg_inf = torch.tensor(float("-inf"), device=s.device, dtype=s.dtype) |
| for _ in range(n_keep): |
| utility = s.clone() |
| pos = s > 0 |
| utility = torch.where(pos, s * (1.0 - r_max), s) |
| utility = torch.where(remaining, utility, neg_inf) |
| j = int(torch.argmax(utility).item()) |
| chosen.append(j) |
| remaining[j] = False |
| r_max = torch.maximum(r_max, r[j]) |
| idx = torch.tensor(chosen, device=s.device, dtype=torch.long) |
| idx, _ = torch.sort(idx) |
| return idx |
|
|
