1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 2327
  6np.random.seed(SEED)
  7random.seed(SEED)
  8
  9
 10def frob(x):
 11    return float(np.linalg.norm(x, 'fro'))
 12
 13
 14def comm(a, b):
 15    return a @ b - b @ a
 16
 17
 18def run_math():
 19    n = 6
 20    sigma = np.arange(n)[::-1]
 21    kappa = np.array([1.0, 1.7, 0.8, 1.3, 2.0, 0.6])
 22    G = np.zeros((n, n))
 23    G[np.arange(n), sigma] = kappa
 24
 25    # A matrix polynomial in G commutes exactly with weighted reflection.
 26    P0 = np.eye(n) + 0.055 * G + 0.009 * (G @ G)
 27    rng = np.random.default_rng(SEED + 1)
 28    E = rng.normal(size=(n, n))
 29    E /= frob(E)
 30    C0 = frob(comm(G, E))
 31
 32    # Prediction 1: ||[G,P0+eps E]|| is exactly linear in eps.
 33    epses = np.array([0.0, .001, .003, .01, .03, .1, .3])
 34    cvals = np.array([frob(comm(G, P0 + e * E)) for e in epses])
 35    slope = float(np.dot(epses, cvals) / max(np.dot(epses, epses), 1e-30))
 36    linear_rel_err = float(np.max(np.abs(cvals - epses * C0)) / max(cvals[-1], 1e-30))
 37
 38    # Prediction 2: exact commutation propagates to powers and resolvents.
 39    power_errors = []
 40    Ppow = np.eye(n)
 41    for t in range(0, 9):
 42        power_errors.append(frob(comm(G, Ppow)))
 43        Ppow = Ppow @ P0
 44    gamma0 = 0.37
 45    R0 = np.linalg.inv(np.eye(n) - (1 - gamma0) * P0)
 46    resolvent_zero = frob(comm(G, R0))
 47
 48    # Prediction 3: [G,R] = (1-gamma) R [G,P] R; it vanishes at gamma=1
 49    # and is approximately linear in alpha=1-gamma near alpha=0.
 50    P1 = P0 + 0.08 * E
 51    gammas = np.array([0.0, .1, .25, .5, .7, .85, .95, 1.0])
 52    rvals, identity_errors = [], []
 53    for gamma in gammas:
 54        alpha = 1.0 - gamma
 55        R = np.linalg.inv(np.eye(n) - alpha * P1)
 56        lhs = comm(G, R)
 57        rhs = alpha * R @ comm(G, P1) @ R
 58        rvals.append(frob(lhs))
 59        identity_errors.append(frob(lhs-rhs) / max(frob(lhs), 1e-30))
 60    rvals = np.array(rvals)
 61    alpha = 1-gammas
 62    near = alpha <= .3
 63    through_origin = float(np.dot(alpha[near], rvals[near]) / max(np.dot(alpha[near], alpha[near]), 1e-30))
 64    # report normalized values so the damping trend is easy to inspect
 65    normalized = (rvals / max(rvals[0], 1e-30)).tolist()
 66
 67    return {
 68        'weighted_G': G.tolist(),
 69        'prediction_1_linear_perturbation': {
 70            'eps': epses.tolist(), 'observed_commutators': cvals.tolist(),
 71            'predicted_slope_norm_[G,E]': C0, 'fit_slope': slope,
 72            'max_relative_deviation': linear_rel_err
 73        },
 74        'prediction_2_exact_multi_step': {
 75            'power_commutator_norms_t0_to_8': power_errors,
 76            'resolvent_commutator_norm_gamma_0.37': resolvent_zero
 77        },
 78        'prediction_3_damping_resolvent': {
 79            'gamma': gammas.tolist(), 'observed_norms': rvals.tolist(),
 80            'normalized_to_gamma_0': normalized,
 81            'resolvent_identity_relative_errors': identity_errors,
 82            'near_gamma_1_fit_slope_vs_(1-gamma)': through_origin,
 83            'gamma_1_norm': float(rvals[-1])
 84        }
 85    }
 86
 87
 88def attention_trial(comm_lambda, seed, steps=180):
 89    import torch
 90    torch.manual_seed(seed); np.random.seed(seed)
 91    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 92    try:
 93        dev = torch.device(device)
 94        N, L, D = 768, 8, 8
 95        x = torch.randn(N, L, D, device=dev)
 96        # Reversal-invariant signal, with deliberately unequal noise scales by side.
 97        latent = x[:, :, 0].sum(1) + .35 * x[:, :, 1].sum(1)
 98        x[:, :L//2] += .45 * torch.randn(N, L//2, D, device=dev)
 99        x[:, L//2:] += 1.25 * torch.randn(N, L//2, D, device=dev)
100        y = (latent > 0).long()
101        tr, va = torch.arange(0, 600, device=dev), torch.arange(600, N, device=dev)
102        q = torch.nn.Linear(D, D, bias=False).to(dev)
103        k = torch.nn.Linear(D, D, bias=False).to(dev)
104        v = torch.nn.Linear(D, D, bias=False).to(dev)
105        out = torch.nn.Linear(D, 2).to(dev)
106        opt = torch.optim.Adam(list(q.parameters())+list(k.parameters())+list(v.parameters())+list(out.parameters()), lr=.012)
107        rev = torch.arange(L-1, -1, -1, device=dev)
108        kap = torch.tensor([1., 1.35, .85, 1.15, 1.15, .85, 1.35, 1.], device=dev)
109        G = torch.zeros(L,L,device=dev); G[torch.arange(L),rev] = kap
110        last_loss = None
111        for _ in range(steps):
112            xb, yb = x[tr], y[tr]
113            P = torch.softmax(q(xb) @ k(xb).transpose(1,2) / math.sqrt(D), dim=-1)
114            z = (P @ v(xb)).mean(1)
115            loss = torch.nn.functional.cross_entropy(out(z), yb)
116            GP = G.unsqueeze(0) @ P; PG = P @ G.unsqueeze(0)
117            penalty = ((GP-PG)**2).sum() / (G.square().sum()*P.square().sum().mean()+1e-8)
118            total = loss + comm_lambda * penalty
119            opt.zero_grad(); total.backward(); opt.step(); last_loss=float(loss.detach())
120        with torch.no_grad():
121            P = torch.softmax(q(x[va]) @ k(x[va]).transpose(1,2) / math.sqrt(D), dim=-1)
122            pred = out((P @ v(x[va])).mean(1)).argmax(1)
123            acc = float((pred == y[va]).float().mean())
124            ce = float(torch.nn.functional.cross_entropy(out((P @ v(x[va])).mean(1)), y[va]))
125            cm = torch.linalg.vector_norm(G@P-P@G).item() / math.sqrt(len(va))
126        return {'accuracy':acc, 'val_loss':ce, 'commutator_rms':cm, 'device':str(dev), 'train_loss':last_loss}
127    except Exception as exc:
128        # CUDA can be shared; retrying on CPU is part of the experiment contract.
129        if device == 'cuda':
130            torch.cuda.empty_cache()
131            old = torch.cuda.is_available
132            # A clean CPU subprocess is simpler than mutating torch's device state.
133            return attention_trial_cpu(comm_lambda, seed, steps)
134        raise
135
136
137def attention_trial_cpu(comm_lambda, seed, steps):
138    import torch
139    # Force CPU by temporarily using a CPU-only equivalent implementation.
140    # The compact fallback duplicates the computation with tensors on CPU.
141    torch.manual_seed(seed); np.random.seed(seed)
142    N,L,D=768,8,8; dev=torch.device('cpu')
143    x=torch.randn(N,L,D); latent=x[:,:,0].sum(1)+.35*x[:,:,1].sum(1)
144    x[:,:4]+=.45*torch.randn(N,4,D); x[:,4:]+=1.25*torch.randn(N,4,D); y=(latent>0).long()
145    q=torch.nn.Linear(D,D,bias=False); k=torch.nn.Linear(D,D,bias=False); v=torch.nn.Linear(D,D,bias=False); out=torch.nn.Linear(D,2)
146    opt=torch.optim.Adam(list(q.parameters())+list(k.parameters())+list(v.parameters())+list(out.parameters()),lr=.012)
147    rev=torch.arange(7,-1,-1); kap=torch.tensor([1.,1.35,.85,1.15,1.15,.85,1.35,1.]); G=torch.zeros(L,L); G[torch.arange(L),rev]=kap
148    for _ in range(steps):
149        P=torch.softmax(q(x[:600])@k(x[:600]).transpose(1,2)/math.sqrt(D),-1); z=(P@v(x[:600])).mean(1); loss=torch.nn.functional.cross_entropy(out(z),y[:600]); pen=((G@P-P@G)**2).sum()/(G.square().sum()*P.square().sum().mean()+1e-8); opt.zero_grad(); (loss+comm_lambda*pen).backward(); opt.step()
150    with torch.no_grad():
151        P=torch.softmax(q(x[600:])@k(x[600:]).transpose(1,2)/math.sqrt(D),-1); logits=out((P@v(x[600:])).mean(1)); return {'accuracy':float((logits.argmax(1)==y[600:]).float().mean()),'val_loss':float(torch.nn.functional.cross_entropy(logits,y[600:])),'commutator_rms':float(torch.linalg.vector_norm(G@P-P@G).item()/math.sqrt(168)),'device':'cpu','train_loss':float(loss)}
152
153
154def main():
155    math_result=run_math()
156    attention={}
157    for lam in (0.0, 1.0):
158        vals=[attention_trial(lam,s) for s in (11,22)]
159        attention[str(lam)]={'seeds':vals,'mean_accuracy':float(np.mean([v['accuracy'] for v in vals])),'mean_val_loss':float(np.mean([v['val_loss'] for v in vals])),'mean_commutator_rms':float(np.mean([v['commutator_rms'] for v in vals]))}
160    result={'seed':SEED,'math':math_result,'attention_mini_experiment':attention}
161    Path('results.json').write_text(json.dumps(result,indent=2))
162    print(json.dumps(result,indent=2))
163
164if __name__=='__main__': main()