diff options
Diffstat (limited to 'src/zbp_scaling/zbp/estimators.py')
| -rw-r--r-- | src/zbp_scaling/zbp/estimators.py | 421 |
1 files changed, 421 insertions, 0 deletions
diff --git a/src/zbp_scaling/zbp/estimators.py b/src/zbp_scaling/zbp/estimators.py new file mode 100644 index 0000000..81f29a3 --- /dev/null +++ b/src/zbp_scaling/zbp/estimators.py @@ -0,0 +1,421 @@ +"""Zeroth-order estimators of g = J^T v for a black-box block. + +Conventions +----------- +A block maps x (shape [*batch_shape, *feat]) to y. A *measurement* function +`h(x_pert)` returns v^T f(x_pert) per measurement unit: + * global units: h -> [*batch_shape] (sum over all feature dims) + * local units: h -> [*batch_shape, H', W'] (sum over channels only; conv blocks) +`h` accepts any integer multiple of the batch stacked along dim 0, so several +probes can be evaluated in one call. + +For local units the directional derivatives D_p (one per output position p) are +box-summed over the receptive field before multiplying the probe, which keeps the +estimator unbiased while removing the variance contributed by far-away positions. + +Every estimator returns (g_hat, queries) where queries is the number of block +evaluations *per sample unit* that a physical system would have to perform. +""" +import math +import torch +import torch.nn.functional as F + +from .probes import sample_probes + + +class Local: + """Receptive-field geometry of a conv-like block: output position p depends on input + positions q with |q - stride*p|_inf <= radius (input coordinates).""" + + def __init__(self, radius, stride=1): + self.radius, self.stride = int(radius), int(stride) + + def gather(self, D, in_hw): + """D: [n, *bs, H', W'] per-output-position derivatives -> S: [n, *bs, H, W] box sums.""" + lead = D.shape[:-2] + Hp, Wp = D.shape[-2:] + H, W = in_hw + k = 2 * self.radius + 1 + Dm = D.reshape(-1, 1, Hp, Wp) + ones = torch.ones(1, 1, k, k, dtype=D.dtype, device=D.device) + oph = H - ((Hp - 1) * self.stride + 1) + opw = W - ((Wp - 1) * self.stride + 1) + S = F.conv_transpose2d(Dm, ones, stride=self.stride, padding=self.radius, output_padding=(oph, opw)) + return S.reshape(*lead, H, W) + + +class Local1D: + """Causal token-window locality for [B, T, d] blocks: output token t depends on input tokens + s in [t-window+1, t]. Measurement units = output tokens; gather = causal window box-sum.""" + kind = "1d" + + def __init__(self, window): + self.window = int(window) + + def gather(self, D, in_shape): + """D: [n, *bs, T] per-output-token derivatives -> S: [n, *bs, T], S[s] = sum_{t=s}^{s+W-1} D[t].""" + T = D.shape[-1] + c = torch.cumsum(D, dim=-1) + pad = torch.zeros_like(c[..., :1]) + c0 = torch.cat([pad, c], dim=-1) # c0[t] = sum_{<t} + hi = torch.clamp(torch.arange(T, device=D.device) + self.window, max=T) + lo = torch.arange(T, device=D.device) + return c0[..., hi] - c0[..., lo] + + +def _probes_for(x, n_probes, family, gen, n_batch_dims, cycle=None): + bs = x.shape[:n_batch_dims] + feat = x.shape[n_batch_dims:] + d = math.prod(feat) + u = sample_probes(n_probes, tuple(bs), d, family, gen, x.device, x.dtype, cycle=cycle) + return u.reshape(n_probes, *bs, *feat), d + + +def zo_vjp_cd_parts(h, x, cfg, gen, n_batch_dims=1, cycle=None): + """Central-difference measurements together with their probes (global units only): + returns (u [n,*bs,*feat], D [n,*bs] = u_i^T g + O(eps^2), g_hat, queries).""" + u, d = _probes_for(x, cfg.n_probes, cfg.probe, gen, n_batch_dims, cycle) + eps = _eps_for(cfg, d) + D = _eval_pairs(h, x, u, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0] + return u, D, _combine(u, D, n_batch_dims, None, x.shape), float(2 * cfg.n_probes) + + +def _combine(u, D, n_batch_dims, local, in_shape): + """g_hat = (1/n) sum_i u_i * S_i, where S_i are the (gathered) directional derivatives.""" + n = u.shape[0] + if local is None: + S = D.reshape(*D.shape, *([1] * (u.dim() - 1 - n_batch_dims))) + elif getattr(local, "kind", "2d") == "1d": + S = local.gather(D, in_shape).unsqueeze(-1) # [n, *bs, T, 1] broadcast over features + else: + S = local.gather(D, in_shape[-2:]) # [n, *bs, H, W] + S = S.unsqueeze(n_batch_dims + 1) # broadcast over channels: [n, *bs, 1, H, W] + return (u * S).mean(0) + + +def _eps_for(cfg, d): + return cfg.eps if cfg.eps_mode == "coord" else cfg.eps / math.sqrt(d) + + +def _eval_pairs(h, x, u, eps_list, n_batch_dims, batch_probes=True, chunk=0): + """D(eps, u_i) = (h(x+eps u_i) - h(x-eps u_i)) / (2 eps) for all probes/eps -> [len(eps), n, *bs, *meas].""" + n = u.shape[0] + B0 = x.shape[0] + outs = [] + if batch_probes: + step = n if chunk <= 0 else max(1, min(chunk, n)) + for eps in eps_list: + parts = [] + for i0 in range(0, n, step): + xs = [] + for i in range(i0, min(n, i0 + step)): + xs.append(x + eps * u[i]) + xs.append(x - eps * u[i]) + hv = h(torch.cat(xs, 0)) + hv = hv.reshape(len(xs) // 2, 2, B0, *hv.shape[1:]) + parts.append((hv[:, 0] - hv[:, 1]) / (2 * eps)) + outs.append(torch.cat(parts, 0)) + else: + for eps in eps_list: + outs.append(torch.stack([(h(x + eps * u[i]) - h(x - eps * u[i])) / (2 * eps) for i in range(n)])) + return torch.stack(outs) + + +def exact_directional(f, x, v, u, n_batch_dims=1, local=None, chunk=16): + """Oracle: exact directional derivatives D_i = <v, J u_i> per measurement unit via forward-mode AD + (vmapped over probes in chunks; falls back to a loop if the block is not vmap-compatible).""" + def one(t): + return torch.func.jvp(f, (x,), (t,))[1] + outs = [] + n = u.shape[0] + try: + for i in range(0, n, chunk): + outs.append(torch.func.vmap(one)(u[i:i + chunk])) + Ju = torch.cat(outs, 0) + except Exception: + Ju = torch.stack([one(u[i]) for i in range(n)]) + prod = Ju * v.unsqueeze(0) + if local is None: + return prod.reshape(n, *x.shape[:n_batch_dims], -1).sum(-1) + if getattr(local, "kind", "2d") == "1d": + return prod.sum(-1) + return prod.sum(n_batch_dims + 1) # sum over channel dim only + + +def oracle_vjp(f, x, v, cfg, gen, n_batch_dims=1, local=None): + """Stage A: (1/n) sum_i u_i <v, J u_i> (== (1/n) sum_i u_i (u_i^T g)). Conditionally unbiased.""" + u, d = _probes_for(x, cfg.n_probes, cfg.probe, gen, n_batch_dims) + with torch.no_grad(): + D = exact_directional(f, x, v, u, n_batch_dims, local) + return _combine(u, D, n_batch_dims, local, x.shape), 0.0 + + +def oracle_projection(g, cfg, gen, n_batch_dims=1): + """Projection form of the oracle for a given exact VJP g (global units only).""" + u, d = _probes_for(g, cfg.n_probes, cfg.probe, gen, n_batch_dims) + n = u.shape[0] + bs = g.shape[:n_batch_dims] + D = (u.reshape(n, *bs, -1) * g.reshape(1, *bs, -1)).sum(-1) + return _combine(u, D, n_batch_dims, None, g.shape), 0.0 + + +def zo_vjp(h, x, cfg, gen, n_batch_dims=1, local=None, cycle=None): + """Function-query estimate of J^T v. Returns (g_hat, queries_per_sample_unit).""" + if cfg.mode == "cd" and cfg.probe == "coordinate_all" and local is None: + # analytic combine: u_i = sqrt(d) e_i -> g_hat[..., i] = D_i / sqrt(d); the probe tensor is a broadcast + # view of eye(d), so memory stays O(d^2 + measurements) instead of O(n * batch * d) + bs = x.shape[:n_batch_dims] + d = math.prod(x.shape[n_batch_dims:]) + if cfg.n_probes != d: + raise ValueError(f"coordinate_all needs n == d (n={cfg.n_probes}, d={d})") + eps = _eps_for(cfg, d) + ident = math.sqrt(d) * torch.eye(d, device=x.device, dtype=x.dtype) + uv = ident.reshape(d, *([1] * len(bs)), d).expand(d, *bs, d) + D = _eval_pairs(h, x, uv, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0] + return (D.movedim(0, -1) / math.sqrt(d)).reshape(x.shape), float(2 * d) + if cfg.mode == "cd" and cfg.probe_chunk and cfg.n_probes > 4 * cfg.probe_chunk and cfg.probe != "hadamard_cycle": + # probe-chunked sample -> evaluate -> accumulate: memory O(chunk * batch * d) for any n + bs = x.shape[:n_batch_dims] + d = math.prod(x.shape[n_batch_dims:]) + eps = _eps_for(cfg, d) + g_acc = torch.zeros_like(x) + done = 0 + while done < cfg.n_probes: + c = min(cfg.probe_chunk, cfg.n_probes - done) + uc, _ = _probes_for(x, c, cfg.probe, gen, n_batch_dims) + D = _eval_pairs(h, x, uc, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0] + g_acc += _combine(uc, D, n_batch_dims, local, x.shape) * c + done += c + return g_acc / cfg.n_probes, float(2 * cfg.n_probes) + u, d = _probes_for(x, cfg.n_probes, cfg.probe, gen, n_batch_dims, cycle) + n = cfg.n_probes + eps = _eps_for(cfg, d) + bs = x.shape[:n_batch_dims] + B0 = x.shape[0] + + if cfg.mode == "forward": + if cfg.batch_probes: + hv = h(torch.cat([x] + [x + eps * u[i] for i in range(n)], 0)) + hv = hv.reshape(n + 1, B0, *hv.shape[1:]) + D = (hv[1:] - hv[:1]) / eps + else: + h0 = h(x) + D = torch.stack([(h(x + eps * u[i]) - h0) / eps for i in range(n)]) + return _combine(u, D, n_batch_dims, local, x.shape), float(n + 1) + + if cfg.mode == "cd": + D = _eval_pairs(h, x, u, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0] + return _combine(u, D, n_batch_dims, local, x.shape), float(2 * n) + + if cfg.mode == "richardson": + D = _eval_pairs(h, x, u, [eps, eps / 2], n_batch_dims, cfg.batch_probes, cfg.probe_chunk) + R = (4 * D[1] - D[0]) / 3 + return _combine(u, R, n_batch_dims, local, x.shape), float(4 * n) + + if cfg.mode in ("ml", "ml_richardson"): + rich = cfg.mode == "ml_richardson" + if cfg.level_sampling == "per_sample": + U = torch.rand((n, *bs), generator=gen, device=x.device, dtype=torch.float64) + else: + U = torch.rand((n,), generator=gen, device=x.device, dtype=torch.float64) + U = U.reshape(n, *([1] * len(bs))).expand(n, *bs) + N = torch.floor(-torch.log2(U) / cfg.alpha).clamp(max=cfg.max_level).to(torch.int64) + if cfg.eps_min > 0: # truncate the ladder where readout noise would dominate (bias O(eps_min^2)) + kmax = max(0, int(math.floor(math.log2(eps / cfg.eps_min)))) if eps > cfg.eps_min else 0 + N = N.clamp(max=kmax) + Nmax = int(N.max().item()) + K = Nmax + (2 if rich else 1) + eps_list = [eps * 2.0 ** (-k) for k in range(K)] + D = _eval_pairs(h, x, u, eps_list, n_batch_dims, cfg.batch_probes, cfg.probe_chunk) # [K, n, *bs, *meas] + S = (4 * D[1:] - D[:-1]) / 3 if rich else D + est = S[0].clone() + extra = S.dim() - 1 - N.dim() + for k in range(1, Nmax + 1): + mask = (N >= k).to(x.dtype).reshape(*N.shape, *([1] * extra)) + est = est + mask * (S[k] - S[k - 1]) / (2.0 ** (-cfg.alpha * k)) + per = 2.0 * (N.to(torch.float64) + 1 + (1 if rich else 0)) + queries = float(per.sum().item() / (math.prod(bs) if len(bs) else 1)) + return _combine(u, est, n_batch_dims, local, x.shape), queries + + raise ValueError(f"unknown ZBP mode {cfg.mode!r}") + + +def causal_score_coordinate_k_vjp(qkv, v_out, cfg, gen, d, n_heads, dk=None, dv=None, window=None): + """Estimate attention's K-side VJP by probing causal score rows. + + The black-box interface here is ``softmax(scores) @ value`` rather than K itself. + One paired query perturbs one valid score coordinate in every (sample, head, + output-token) row, and the output supplies a separate scalar measurement for + every such row. For a row with ``m`` valid keys, the estimator visits all + coordinates when ``m <= n``; otherwise it samples ``n`` coordinates without + replacement and applies the Horvitz--Thompson weight ``m / n``. Thus the + row-gradient relative MSE is exactly ``max(m / n - 1, 0)`` in the directional- + derivative limit, instead of ``(m - 1) / n`` for i.i.d. dense probes. + + The estimated score VJP is mapped to K through the digitally known bilinear + score map QK^T/sqrt(dh). Q and V are held fixed during the function queries. + Returns ``(g_k, queries_per_sequence)``; full-attention K is [B,T,d], while + multi-query K is [B,T,dh]. + """ + if cfg.mode not in ("cd", "oracle"): + raise ValueError(f"causal score-coordinate probing supports cd/oracle, got {cfg.mode!r}") + B, T, _ = qkv.shape + H = int(n_heads) + dh = int(d) // H + dk = int(d if dk is None else dk) + dv = int(d if dv is None else dv) + if int(d) % H: + raise ValueError(f"attention width {d} is not divisible by {H} heads") + if dk not in (int(d), dh) or dv not in (int(d), dh): + raise ValueError("score-coordinate probing supports full or multi-query K/V layouts") + + q0, k0, value0 = qkv.split([int(d), dk, dv], dim=-1) + q = q0.view(B, T, H, dh).transpose(1, 2) + if dk == dh: + k = k0.view(B, T, 1, dh).transpose(1, 2).expand(B, H, T, dh) + else: + k = k0.view(B, T, H, dh).transpose(1, 2) + if dv == dh: + value = value0.view(B, T, 1, dh).transpose(1, 2).expand(B, H, T, dh) + else: + value = value0.view(B, T, H, dh).transpose(1, 2) + delta = v_out.view(B, T, H, dh).transpose(1, 2) + + idx = torch.arange(T, device=qkv.device) + valid = idx[None, :] <= idx[:, None] + if window is not None: + valid = valid & (idx[None, :] > idx[:, None] - int(window)) + counts = valid.sum(-1) # [T], valid coordinates per row + scores = (q @ k.transpose(-2, -1)) / math.sqrt(dh) + scores = scores.masked_fill(~valid, float("-inf")) + + # Analytic score VJP is useful for projection-only audits. The CD path below + # obtains the same quantities solely from paired softmax/value evaluations. + if cfg.mode == "oracle": + att = scores.softmax(-1) + g_att = delta @ value.transpose(-2, -1) + g_score_exact = att * (g_att - (att * g_att).sum(-1, keepdim=True)) + g_score_exact = g_score_exact.masked_fill(~valid, 0) + + n = int(cfg.n_probes) + if n <= 0: + raise ValueError("score-coordinate probing needs at least one probe") + slots = min(n, T) + # Independent random subsets for every sample/head/row. Invalid positions + # sort last, so the first min(n,m) entries are a uniform subset of the m-prefix. + priorities = torch.rand((B, H, T, T), generator=gen, device=qkv.device, dtype=torch.float32) + priorities.masked_fill_(~valid.view(1, 1, T, T), 2.0) + chosen = priorities.argsort(dim=-1)[..., :slots] # [B,H,T,slots] + chosen = chosen.permute(3, 0, 1, 2).contiguous() # [slots,B,H,T] + slot_id = torch.arange(slots, device=qkv.device).view(slots, 1, 1, 1) + active = slot_id < counts.view(1, 1, 1, T) + active = active.expand(slots, B, H, T) + weights = torch.where(counts <= n, torch.ones_like(counts, dtype=qkv.dtype), + counts.to(qkv.dtype) / n) # [T] + g_score = torch.zeros_like(scores) + + if cfg.mode == "oracle": + picked = torch.gather(g_score_exact.unsqueeze(0).expand(slots, -1, -1, -1, -1), + -1, chosen.unsqueeze(-1)).squeeze(-1) + contrib = picked * active.to(qkv.dtype) * weights.view(1, 1, 1, T) + for j in range(slots): + g_score.scatter_add_(-1, chosen[j].unsqueeze(-1), contrib[j].unsqueeze(-1)) + queries = 0.0 + else: + eps = float(cfg.eps) # directions are unit coordinates, so coord/norm conventions coincide + if eps <= 0: + raise ValueError("score-coordinate probing needs eps > 0") + chunk = slots if cfg.probe_chunk <= 0 else max(1, min(int(cfg.probe_chunk), slots)) + for i0 in range(0, slots, chunk): + ids = chosen[i0:i0 + chunk] + act = active[i0:i0 + chunk] + c = ids.shape[0] + plus = scores.unsqueeze(0).expand(c, -1, -1, -1, -1).clone() + minus = plus.clone() + shift = act.to(qkv.dtype).unsqueeze(-1) * eps + plus.scatter_add_(-1, ids.unsqueeze(-1), shift) + minus.scatter_add_(-1, ids.unsqueeze(-1), -shift) + y_plus = plus.softmax(-1) @ value.unsqueeze(0) + y_minus = minus.softmax(-1) @ value.unsqueeze(0) + if cfg.readout_noise > 0: + y_plus = y_plus + cfg.readout_noise * torch.randn( + y_plus.shape, generator=gen, device=y_plus.device, dtype=y_plus.dtype) + y_minus = y_minus + cfg.readout_noise * torch.randn( + y_minus.shape, generator=gen, device=y_minus.device, dtype=y_minus.dtype) + deriv = ((y_plus - y_minus) * delta.unsqueeze(0)).sum(-1) / (2 * eps) + contrib = deriv * act.to(qkv.dtype) * weights.view(1, 1, 1, T) + for j in range(c): + g_score.scatter_add_(-1, ids[j].unsqueeze(-1), contrib[j].unsqueeze(-1)) + queries = float(2 * n) + + _LAST_SCORE_VJP['g'] = g_score + gk_heads = g_score.transpose(-2, -1) @ q / math.sqrt(dh) + if dk == dh: + gk = gk_heads.sum(1).transpose(1, 2) # shared K accumulates all heads + else: + gk = gk_heads.transpose(1, 2).reshape(B, T, int(d)) + return gk, queries + + +def make_h(f, v, n_batch_dims=1, local=None, readout_noise=0.0, gen=None): + """Measurement function h(x_pert) = <v, f(x_pert)> per unit, tolerant to stacked batches. + readout_noise > 0 adds i.i.d. Gaussian noise to every measured output coordinate of every query + (a physical readout-noise model; v is applied digitally afterwards).""" + B0 = v.shape[0] + + def h(xp): + with torch.no_grad(): + y = f(xp) + if readout_noise > 0: + y = y + readout_noise * torch.randn(y.shape, generator=gen, device=y.device, dtype=y.dtype) + k = y.shape[0] // B0 + vv = v.repeat(k, *([1] * (v.dim() - 1))) if k > 1 else v + prod = y * vv + if local is None: + return prod.reshape(*prod.shape[:n_batch_dims], -1).sum(-1) + if getattr(local, "kind", "2d") == "1d": + return prod.sum(-1) # per output token + return prod.sum(n_batch_dims) # keep spatial dims + return h + + +def control_vjp(g, cfg, gen, n_batch_dims=1): + """Control estimators. 'zero': no credit is propagated (downstream layers frozen). + 'noise': isotropic Gaussian noise with the per-sample norm that the oracle estimate with + the same n would have (sqrt(1 + (d-1)/n) ||g||): same magnitude, zero information.""" + if cfg.mode == "zero": + return torch.zeros_like(g), 0.0 + if cfg.mode == "noise": + bs = g.shape[:n_batch_dims] + gf = g.reshape(*bs, -1) + d = gf.shape[-1] + z = torch.randn(gf.shape, generator=gen, device=g.device, dtype=g.dtype) + z = z / z.norm(dim=-1, keepdim=True).clamp_min(1e-30) + target = gf.norm(dim=-1, keepdim=True) * math.sqrt(1 + (d - 1) / cfg.n_probes) + return (z * target).reshape(g.shape), 0.0 + raise ValueError(cfg.mode) + + +def causal_score_coordinate_qk_vjp(qkv, v_out, cfg, gen, d, n_heads, dk=None, dv=None, window=None): + """Same score-space probing as causal_score_coordinate_k_vjp, but returns the Q-side VJP as well + (both follow from the score-VJP through the known bilinear map S = Q K^T / sqrt(d_h)).""" + B, T, _ = qkv.shape + H = int(n_heads); dh = int(d) // H + dk = int(d if dk is None else dk) + gk, queries = causal_score_coordinate_k_vjp(qkv, v_out, cfg, gen, d, n_heads, dk, dv, window) + # recover g_score from gk is not possible in general; recompute the score VJP estimate via the same + # probes by calling the K routine on a transposed view is not exact either, so instead re-derive: + # g_K[s] = sum_t g_S[t,s] q_t / sqrt(dh) and g_Q[t] = sum_s g_S[t,s] k_s / sqrt(dh). + # We obtain g_S by solving nothing: the K routine exposes it through `_last_score_vjp`. + g_score = _LAST_SCORE_VJP["g"] + k0 = qkv[..., int(d):int(d) + dk] + if dk == dh: + k = k0.view(B, T, 1, dh).transpose(1, 2).expand(B, H, T, dh) + else: + k = k0.view(B, T, H, dh).transpose(1, 2) + gq = (g_score @ k) / math.sqrt(dh) # [B,H,T,dh] + gq = gq.transpose(1, 2).reshape(B, T, int(d)) + return gq, gk, queries + + +_LAST_SCORE_VJP = {"g": None} |