1"""Moment-Sharp spectral control for Linear layers.
 2
 3This module is a training intervention, not a replacement for the mandated bench
 4harness.  K=2 uses the exact nonnegative moment extremizer and is differentiable
 5only through the post-step rescaling operation (the controller itself is kept
 6outside autograd).
 7"""
 8import math
 9import torch
10from torch import nn
11
12
13def u2_from_moments(m1, m2, d, eps=1e-12):
14    """Maximum possible eigenvalue given sum and sum of squares."""
15    d = int(d)
16    if d <= 1:
17        return max(float(m1), eps)
18    disc = max(0.0, d * float(m2) - float(m1) ** 2)
19    return max(float(m1) / d + math.sqrt((d - 1) * disc) / d, eps)
20
21
22def exact_or_hutchinson_moments(weight, probes=8, generator=None):
23    """Estimate tr(A), tr(A^2), A=W^T W, with exact small-matrix fallback.
24
25    Hutchinson is used for larger matrices to keep the mechanism aligned with
26    the proposed implementation; returned values are detached Python floats.
27    """
28    w = weight.detach()
29    d = w.shape[1]
30    if d <= 256:
31        a = w.T @ w
32        return float(torch.trace(a)), float(torch.trace(a @ a)), d
33    gen = generator or torch.Generator(device=w.device)
34    if generator is None:
35        gen.manual_seed(12345)
36    z = torch.randint(0, 2, (probes, d), device=w.device, generator=gen,
37                      dtype=torch.int64).to(w.dtype).mul_(2).sub_(1)
38    az = (z @ w.T) @ w
39    a2z = (az @ w.T) @ w
40    return float((z * az).sum(1).mean()), float((z * a2z).sum(1).mean()), d
41
42
43@torch.no_grad()
44def moment_sharp_rescale(model, target_sigma=2.0, probes=8, ema=None,
45                         generator=None):
46    """Rescale each Linear weight when its K=2 certified bound exceeds target."""
47    observed = []
48    for layer in model.modules():
49        if not isinstance(layer, nn.Linear):
50            continue
51        m1, m2, d = exact_or_hutchinson_moments(layer.weight, probes, generator)
52        key = id(layer)
53        if ema is not None:
54            old = ema.get(key, (m1, m2))
55            m1, m2 = .9 * old[0] + .1 * m1, .9 * old[1] + .1 * m2
56            ema[key] = (m1, m2)
57        bound_sq = u2_from_moments(m1, m2, d)
58        before = float(torch.linalg.matrix_norm(layer.weight, 2))
59        if bound_sq > target_sigma ** 2:
60            layer.weight.mul_(target_sigma / math.sqrt(bound_sq + 1e-12))
61        after = float(torch.linalg.matrix_norm(layer.weight, 2))
62        observed.append({"bound_sigma": math.sqrt(bound_sq),
63                         "true_sigma_before": before,
64                         "true_sigma_after": after})
65    return observed
66
67
68def mechanism_signature(model, target_sigma=2.0):
69    """Measured NN-scale signature: bound, observed sigma, and certified status."""
70    rows = []
71    for layer in model.modules():
72        if isinstance(layer, nn.Linear):
73            m1, m2, d = exact_or_hutchinson_moments(layer.weight, probes=32)
74            b = math.sqrt(u2_from_moments(m1, m2, d))
75            s = float(torch.linalg.matrix_norm(layer.weight.detach(), 2))
76            rows.append((b, s))
77    return {"predicted_bound_sigma": [x[0] for x in rows],
78            "observed_true_sigma": [x[1] for x in rows],
79            "max_bound_minus_observed": max((b-s for b,s in rows), default=0.),
80            "target_sigma": target_sigma,
81            "confirmed": bool(all(b + 1e-5 >= s for b,s in rows))}