Rank-One Feedback Spectrum Regularizer / experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
 1import json, math, random
 2import numpy as np
 3import torch
 4
 5SEED = 381
 6np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
 7torch.set_num_threads(4)
 8
 9
10def determinant_lemma_check():
11    n = 6
12    R = np.random.randn(n, n)
13    A = 0.65 * R / max(abs(np.linalg.eigvals(R)))
14    b, c = np.random.randn(n), np.random.randn(n)
15    J = A + np.outer(b, c)
16    pts = [.7+.4j, -1.1+.2j, .3-.9j]
17    errors = []
18    for z in pts:
19        lhs = np.linalg.det(z*np.eye(n)-J)
20        q = np.linalg.solve(z*np.eye(n)-A, b)
21        rhs = np.linalg.det(z*np.eye(n)-A) * (1-c@q)
22        errors.append(abs(lhs-rhs)/max(1., abs(lhs), abs(rhs)))
23    eigJ = np.linalg.eigvals(J); eigA = np.linalg.eigvals(A)
24    mode_errors = []
25    for z in eigJ:
26        if min(abs(z-eigA)) > 1e-6:
27            mode_errors.append(abs(1-c@np.linalg.solve(z*np.eye(n)-A,b)))
28    return {"determinant_relative_error": float(max(errors)),
29            "mode_characteristic_error": float(max(mode_errors)),
30            "open_loop_radius": float(max(abs(eigA))),
31            "closed_loop_radius": float(max(abs(eigJ)))}
32
33
34def contour_penalty(g, a=.92, margin=.10, radii=(.98, 1.0, 1.02), ntheta=16):
35    # A=aI, b=c=sqrt(g), hence H(z)=g/(z-a); this is still evaluated by a solve.
36    A = torch.eye(1, dtype=torch.complex64) * a
37    b = torch.sqrt(torch.clamp(g, min=1e-8)).to(torch.complex64).reshape(1)
38    c = b
39    vals = []
40    for r in radii:
41        for k in range(ntheta):
42            th = 2*math.pi*k/ntheta
43            z = torch.tensor(r*math.cos(th)+1j*r*math.sin(th), dtype=torch.complex64)
44            q = torch.linalg.solve(z*torch.eye(1, dtype=torch.complex64)-A, b)
45            vals.append(torch.relu(torch.abs(torch.sum(c*q))-(1-margin))**2)
46    return torch.stack(vals).mean()
47
48
49def train(reg, steps=240, lr=.035):
50    # The task optimum deliberately asks for a strong positive feedback gain.
51    # Regularization should prevent the associated pole a+g from crossing one.
52    g = torch.tensor(0.02, requires_grad=True)
53    opt = torch.optim.SGD([g], lr=lr)
54    trace=[]
55    for step in range(steps):
56        opt.zero_grad()
57        task = (g-0.18)**2
58        pen = contour_penalty(g) if reg else 0*g
59        loss = task + 1.8*pen
60        loss.backward(); opt.step()
61        with torch.no_grad(): g.clamp_(0, .35)
62        if step % 20 == 0 or step == steps-1:
63            trace.append((step, float(g), float(task), float(pen)))
64    gv=float(g)
65    pole=.92+gv
66    # impulse response of the scalar closed-loop recurrence
67    impulse=[pole**k for k in range(20)]
68    return {"g":gv, "task_mse":(gv-.18)**2, "max_contour_gain":max(
69        float(torch.abs(torch.tensor(gv/(r*math.cos(2*math.pi*k/16)+1j*r*math.sin(2*math.pi*k/16)-.92))))
70        for r in (.98,1.,1.02) for k in range(16)),
71        "closed_loop_pole":pole, "impulse_abs_k19":abs(impulse[-1]), "trace":trace}
72
73
74def main():
75    check=determinant_lemma_check()
76    base=train(False); idea=train(True)
77    result={"seed":SEED, "math_check":check, "baseline":base, "idea":idea,
78            "claim_signal": {"lower_pole": idea["closed_loop_pole"] < base["closed_loop_pole"],
79                             "lower_contour_gain": idea["max_contour_gain"] < base["max_contour_gain"],
80                             "faster_impulse_decay": idea["impulse_abs_k19"] < base["impulse_abs_k19"]}}
81    with open("results.json","w") as f: json.dump(result,f,indent=2)
82    print(json.dumps(result,indent=2))
83
84if __name__ == "__main__": main()