1import json
 2import math
 3import random
 4import numpy as np
 5import torch
 6
 7# CPU-scale, deterministic benchmark for normalized class prototypes.
 8torch.set_num_threads(4)
 9SEED = 123
10
11def seed_all(seed):
12    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
13
14def normalize(x):
15    return x / x.norm(dim=1, keepdim=True).clamp_min(1e-12)
16
17def math_check():
18    seed_all(SEED)
19    n, d = 32, 24
20    x = normalize(torch.randn(n, d))
21    g = torch.randn(n, d)
22    gx = g - (g*x).sum(1, keepdim=True)*x
23    v = torch.randn(n, d)
24    v = v - (v*x).sum(1, keepdim=True)*x
25    eta = 0.17
26    trial = normalize(x + eta*v)
27    v2 = v - (v*trial).sum(1, keepdim=True)*trial
28    # P is symmetric/idempotent, projected vectors are tangent, and retraction is unit norm.
29    px = g - (g*x).sum(1, keepdim=True)*x
30    return {
31        "projector_idempotence": float((px - (px - (px*x).sum(1, keepdim=True)*x)).abs().max()),
32        "projected_tangent_max": float((gx*x).sum(1).abs().max()),
33        "retraction_norm_error": float((trial.norm(dim=1)-1).abs().max()),
34        "velocity_tangent_max": float((v2*trial).sum(1).abs().max()),
35    }
36
37def make_data(seed, classes=20, d=24, train_per=80, test_per=40):
38    gen = torch.Generator().manual_seed(seed)
39    centers = normalize(torch.randn(classes, d, generator=gen))
40    def sample(per, noise):
41        y = torch.arange(classes).repeat_interleave(per)
42        z = centers[y] + noise*torch.randn(len(y), d, generator=gen)
43        return normalize(z), y
44    return centers, sample(train_per, .55), sample(test_per, .55)
45
46def run(method, seed, steps=220, lr=.28, beta=.9, scale=12.0):
47    centers, (train_x, train_y), (test_x, test_y) = make_data(seed)
48    # Same initialization and data for both methods.
49    gen = torch.Generator().manual_seed(seed+99)
50    p0 = normalize(torch.randn(20, 24, generator=gen))
51    p = p0.clone()
52    v = torch.zeros_like(p)
53    losses=[]; accs=[]; norm_err=[]; tangent=[]
54    for step in range(steps):
55        logits = scale * train_x @ p.T
56        # Full batch makes optimizer differences easy to interpret and reproducible.
57        loss = torch.nn.functional.cross_entropy(logits, train_y)
58        # Make prototypes leaf tensors so gradients are computed independently each step.
59        p = p.detach().requires_grad_(True)
60        loss = torch.nn.functional.cross_entropy(scale * train_x @ p.T, train_y)
61        g, = torch.autograd.grad(loss, p)
62        p = p.detach()
63        if method == "baseline":
64            # Conventional momentum in ambient coordinates, followed by ad hoc normalization.
65            v = beta*v - lr*g
66            p = normalize(p + v)
67        else:
68            # Riemannian inertial update: project force, integrate tangent velocity,
69            # retract, then remove radial velocity component at the new point.
70            gp = g - (g*p).sum(1, keepdim=True)*p
71            v = beta*v - lr*gp
72            p = normalize(p + v)
73            v = v - (v*p).sum(1, keepdim=True)*p
74        with torch.no_grad():
75            te = torch.nn.functional.cross_entropy(scale * test_x @ p.T, test_y)
76            ac = (scale * test_x @ p.T).argmax(1).eq(test_y).float().mean()
77            losses.append(float(loss)); accs.append(float(ac)); norm_err.append(float((p.norm(dim=1)-1).abs().max()))
78            tangent.append(float((v*p).sum(1).abs().max()))
79    return {"final_train_loss": losses[-1], "final_test_loss": float(te),
80            "final_accuracy": accs[-1], "best_accuracy": max(accs),
81            "step_to_70pct": next((i+1 for i,a in enumerate(accs) if a>=.70), None),
82            "norm_error_max": max(norm_err), "velocity_tangent_max": max(tangent),
83            "loss_curve": losses, "accuracy_curve": accs}
84
85def main():
86    check = math_check()
87    all_results={"math_check":check, "runs":{}}
88    for method in ["baseline", "idea"]:
89        rs=[run(method, s) for s in [10,20,30]]
90        def avg(k):
91            vals=[r[k] for r in rs if r[k] is not None]
92            return float(np.mean(vals)) if vals else None
93        all_results["runs"][method]={k:avg(k) for k in ["final_train_loss","final_test_loss","final_accuracy","best_accuracy","step_to_70pct","norm_error_max","velocity_tangent_max"]}
94        all_results["runs"][method]["per_seed"]= [{k:r[k] for k in ["final_train_loss","final_accuracy","step_to_70pct"]} for r in rs]
95    print(json.dumps(all_results, indent=2))
96
97if __name__ == '__main__': main()