Path-Holonomy Attention / holonomy_experiment.py

✓✓ Beats tuned baseline

Raw ⬇ ZIP
 1import json, random
 2import numpy as np
 3import torch
 4from torch import nn
 5
 6SEED = 234
 7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 8
 9def math_check():
10    # Two invertible, noncommuting edge operators.
11    A = np.array([[1., 1.], [0., 1.]])
12    B = np.array([[1., 0.], [1., 1.]])
13    ab, ba = A @ B, B @ A
14    noncomm = float(np.linalg.norm(ab - ba))
15    # Composition law for identity sigma/tau is ordinary ordered multiplication.
16    C = np.array([[2., 0.], [0., .5]])
17    direct = A @ B @ C
18    composed = (A @ B) @ C
19    composition_error = float(np.linalg.norm(direct - composed))
20    # Same additive aggregate, different ordered holonomy.
21    additive_collision = float(np.linalg.norm((A+B) - (B+A)))
22    return {"AB_minus_BA_fro": noncomm, "composition_error": composition_error,
23            "additive_collision": additive_collision, "AB": ab.tolist(), "BA": ba.tolist()}
24
25class Additive(nn.Module):
26    def __init__(self, d=3):
27        super().__init__(); self.edge = nn.Parameter(torch.randn(2,d,d)*.35)
28        self.head = nn.Sequential(nn.Flatten(), nn.Linear(d*d, 16), nn.Tanh(), nn.Linear(16,1))
29    def forward(self, x):
30        return self.head(self.edge[x].sum(1)).squeeze(-1)
31
32class Holonomy(nn.Module):
33    def __init__(self, d=3):
34        super().__init__(); self.edge = nn.Parameter(torch.randn(2,d,d)*.35)
35        # Learned reversal and color-switch maps, implemented as conjugations.
36        self.R = nn.Parameter(torch.eye(d) + .03*torch.randn(d,d))
37        self.S = nn.Parameter(torch.eye(d) + .03*torch.randn(d,d))
38        self.head = nn.Sequential(nn.Flatten(), nn.Linear(d*d, 16), nn.Tanh(), nn.Linear(16,1))
39    def transform(self, X, step):
40        # T is reversal; Sigma is a learned color switch. Alternation follows T^r.
41        Ri = torch.linalg.pinv(self.R)
42        X = self.R @ X @ Ri if step % 2 else X
43        if step % 2: # deterministic two-color channel for this toy path
44            Si = torch.linalg.pinv(self.S); X = Si @ X @ self.S
45        return X
46    def forward(self, x):
47        B,L = x.shape; d=self.edge.shape[-1]
48        H = torch.eye(d, device=x.device).expand(B,d,d).clone()
49        for r in range(L): H = H @ self.transform(self.edge[x[:,r]], r)
50        return self.head(H).squeeze(-1)
51    def regularizer(self):
52        d=self.S.shape[0]; I=torch.eye(d,device=self.S.device)
53        return ((self.S@self.S-I)**2).mean()
54
55def make_data(n, seed):
56    g=np.random.default_rng(seed); x=g.integers(0,2,size=(n,3),dtype=np.int64)
57    # Ordered relation: first two edges must be 0 then 1. Same-count permutations conflict.
58    y=((x[:,0]==0)&(x[:,1]==1)).astype(np.float32)
59    return torch.tensor(x), torch.tensor(y)
60
61def train(model, xt, yt, xv, yv, epochs=450):
62    opt=torch.optim.Adam(model.parameters(),lr=.025,weight_decay=1e-4)
63    lossfn=nn.BCEWithLogitsLoss()
64    for _ in range(epochs):
65        opt.zero_grad(); z=model(xt); loss=lossfn(z,yt)
66        if isinstance(model,Holonomy): loss=loss + .01*model.regularizer()
67        loss.backward(); opt.step()
68    with torch.no_grad():
69        pred=(torch.sigmoid(model(xv))>.5).float(); acc=float((pred==yv).float().mean())
70        train_acc=float(((torch.sigmoid(model(xt))>.5).float()==yt).float().mean())
71    return train_acc,acc
72
73def main():
74    device='cuda' if torch.cuda.is_available() else 'cpu'
75    try:
76        xt,yt=make_data(2048,10); xv,yv=make_data(2048,11)
77        xt,yt,xv,yv=[z.to(device) for z in (xt,yt,xv,yv)]
78        # identical initialization scale and exact same data for fair comparison
79        torch.manual_seed(SEED); base=Additive().to(device)
80        torch.manual_seed(SEED); idea=Holonomy().to(device)
81        ba, bv=train(base,xt,yt,xv,yv); ia,iv=train(idea,xt,yt,xv,yv)
82        result={'device':device,'math':math_check(),
83                'baseline':{'train_accuracy':ba,'test_accuracy':bv},
84                'idea':{'train_accuracy':ia,'test_accuracy':iv},
85                'n_train':len(xt),'n_test':len(xv),'epochs':450}
86    except Exception as e:
87        if device=='cuda':
88            torch.cuda.empty_cache(); device='cpu'
89            xt,yt=make_data(2048,10); xv,yv=make_data(2048,11)
90            torch.manual_seed(SEED); ba,bv=train(Additive(),xt,yt,xv,yv)
91            torch.manual_seed(SEED); ia,iv=train(Holonomy(),xt,yt,xv,yv)
92            result={'device':device,'math':math_check(),'baseline':{'train_accuracy':ba,'test_accuracy':bv},'idea':{'train_accuracy':ia,'test_accuracy':iv},'fallback_reason':str(e)}
93        else: raise
94    print(json.dumps(result,sort_keys=True))
95
96if __name__=='__main__': main()