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
|
"""Bias / variance diagnostics for gradient estimators."""
import math
import torch
def cosine(a, b):
a = a.reshape(a.shape[0], -1) if a.dim() > 1 else a.reshape(1, -1)
b = b.reshape(b.shape[0], -1) if b.dim() > 1 else b.reshape(1, -1)
return torch.nn.functional.cosine_similarity(a, b, dim=1)
def relerr(a, b):
a = a.reshape(a.shape[0], -1)
b = b.reshape(b.shape[0], -1)
return (a - b).norm(dim=1) / b.norm(dim=1).clamp_min(1e-30)
def audit(samples, g):
"""samples: [M, *shape] independent estimates; g: exact [*shape].
Returns dict with relative bias, variance ratio, a chi-square test of zero bias,
and the distribution of cosines.
"""
M = samples.shape[0]
S = samples.reshape(M, -1).double()
g = g.reshape(-1).double()
mean = S.mean(0)
bias = mean - g
var = S.var(0, unbiased=True) # per-coordinate variance
gn2 = g.dot(g).clamp_min(1e-300)
rel_bias = (bias.norm() / gn2.sqrt()).item()
var_ratio = (var.sum() / gn2).item() # E||g_hat - mean||^2 / ||g||^2
# chi-square statistic of H0: bias = 0, per coordinate z_j = bias_j / (sd_j / sqrt(M)),
# restricted to coordinates that are actually random (exact coordinates have zero variance)
active = var > 1e-12 * var.max().clamp_min(1e-300)
se = (var[active] / M).sqrt().clamp_min(1e-300)
z = bias[active] / se
d = max(int(z.numel()), 1)
chi2 = (z * z).sum().item()
# under H0 each z_j is Student-t with nu = M-1 dof: E[t^2] = nu/(nu-2), Var[t^2] = 2 nu^2 (nu-1) / ((nu-2)^2 (nu-4))
nu = M - 1
if nu > 4:
mu2 = nu / (nu - 2)
v2 = 2 * nu ** 2 * (nu - 1) / ((nu - 2) ** 2 * (nu - 4))
else:
mu2, v2 = 1.0, 2.0
zscore = (chi2 - d * mu2) / math.sqrt(d * v2) # approx N(0,1) under H0 for large d
# relative bias expected under H0 (pure noise): sqrt(sum var / M) / ||g||
rel_bias_null = ((var.sum() / M).sqrt() / gn2.sqrt()).item()
cos = torch.nn.functional.cosine_similarity(S, g[None], dim=1)
return {
"M": M, "d": d, "d_total": g.numel(),
"rel_bias": rel_bias,
"rel_bias_null": rel_bias_null,
"bias_ratio": rel_bias / max(rel_bias_null, 1e-300),
"var_ratio": var_ratio,
"chi2": chi2, "chi2_z": zscore,
"cos_mean": cos.mean().item(), "cos_std": cos.std().item(),
"cos_q10": cos.quantile(0.1).item(), "cos_q90": cos.quantile(0.9).item(),
"cos_of_mean": torch.nn.functional.cosine_similarity(mean[None], g[None], dim=1).item(),
}
|
