File size: 2,861 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 | """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
|