1import json, math, time
  2import numpy as np
  3import torch
  4from gated_attention import local_attention, global_linear_attention, dense_attention
  5
  6torch.manual_seed(7)
  7np.random.seed(7)
  8torch.set_num_threads(4)
  9device = "cuda" if torch.cuda.is_available() else "cpu"
 10
 11def graph(n, degree):
 12    src, dst = [], []
 13    for i in range(n):
 14        for t in range(degree):
 15            src.append(i); dst.append((i + t) % n)
 16    return torch.tensor(src, device=device), torch.tensor(dst, device=device)
 17
 18def kernel_explicit(x, wq, wk, wv):
 19    q, k, v = x @ wq, x @ wk, x @ wv
 20    pq, pk = torch.nn.functional.elu(q) + 1, torch.nn.functional.elu(k) + 1
 21    weights = pq @ pk.T
 22    return (weights @ v) / (weights.sum(1, keepdim=True) + 1e-6)
 23
 24def timed(fn, repeats=8):
 25    if device == "cuda": torch.cuda.synchronize()
 26    for _ in range(2): fn()
 27    if device == "cuda": torch.cuda.synchronize()
 28    t0 = time.perf_counter()
 29    for _ in range(repeats): fn()
 30    if device == "cuda": torch.cuda.synchronize()
 31    return (time.perf_counter() - t0) / repeats * 1000
 32
 33# Prediction 1: local attention is exactly dense attention with non-neighbors removed.
 34n, d, r = 31, 12, 7
 35x = torch.randn(n, d, device=device)
 36wq, wk, wv = [torch.randn(d, r, device=device) / math.sqrt(d) for _ in range(3)]
 37src, dst = graph(n, 4)
 38local = local_attention(x, src, dst, wq, wk, wv)
 39q, k, v = x @ wq, x @ wk, x @ wv
 40logits = torch.full((n, n), -torch.inf, device=device)
 41logits[src, dst] = (q[src] * k[dst]).sum(-1) / math.sqrt(r)
 42masked = torch.softmax(logits, dim=-1) @ v
 43mask_error = (local - masked).abs().max().item()
 44
 45# Prediction 2: feature-map linear attention is the exact reassociation of its
 46# pairwise kernel, up to the stated epsilon in the denominator.
 47explicit = kernel_explicit(x, wq, wk, wv)
 48linear = global_linear_attention(x, wq, wk, wv)
 49linear_error = (explicit - linear).abs().max().item()
 50
 51# Prediction 3: operation proxy has dense N^2*r versus E*r + N*r^2;
 52# for fixed degree, dense/hybrid proxy ratio grows linearly with N.
 53width = 16
 54sizes = [64, 128, 256, 512]
 55degree = 8
 56ratios, slopes = [], []
 57for nn in sizes:
 58    e = nn * degree
 59    dense_ops = nn * nn * width
 60    hybrid_ops = e * width + nn * width * width
 61    ratios.append(dense_ops / hybrid_ops)
 62# log-log slope of ratio versus N; predicted slope approaches 1 when N >> r.
 63slope = float(np.polyfit(np.log(sizes), np.log(ratios), 1)[0])
 64
 65# Small learning task: target combines one-hop neighbor mean and global mean.
 66# Dense baseline uses full attention; hybrid uses local+global branches with a
 67# learned gate. Both have the same projections and output head dimensions.
 68def make_data(nn=48, dd=10, deg=5, batches=24):
 69    xx, yy = [], []
 70    ss, tt = graph(nn, deg)
 71    for _ in range(batches):
 72        z = torch.randn(nn, dd, device=device)
 73        neigh = torch.zeros_like(z)
 74        neigh.index_add_(0, ss, z[tt])
 75        counts = torch.zeros(nn, device=device)
 76        counts.index_add_(0, ss, torch.ones_like(ss, dtype=torch.float))
 77        neigh = neigh / counts[:, None]
 78        target = 0.65 * neigh[:, :1] + 0.35 * z.mean(0, keepdim=True)[:, :1]
 79        xx.append(z); yy.append(target)
 80    return xx, yy, ss, tt
 81
 82class DenseModel(torch.nn.Module):
 83    def __init__(self, dd=10, rr=8):
 84        super().__init__(); self.q=torch.nn.Linear(dd,rr,bias=False); self.k=torch.nn.Linear(dd,rr,bias=False); self.v=torch.nn.Linear(dd,rr,bias=False); self.head=torch.nn.Linear(rr,1)
 85    def forward(self,x): return self.head(dense_attention(x,self.q.weight.T,self.k.weight.T,self.v.weight.T))
 86
 87class HybridModel(torch.nn.Module):
 88    def __init__(self, dd=10, rr=8):
 89        super().__init__(); self.q=torch.nn.Linear(dd,rr,bias=False); self.k=torch.nn.Linear(dd,rr,bias=False); self.v=torch.nn.Linear(dd,rr,bias=False); self.g=torch.nn.Linear(dd,1); self.head=torch.nn.Linear(rr,1)
 90    def forward(self,x,ss,tt):
 91        l=local_attention(x,ss,tt,self.q.weight.T,self.k.weight.T,self.v.weight.T); g=global_linear_attention(x,self.q.weight.T,self.k.weight.T,self.v.weight.T); gate=torch.sigmoid(self.g(x)); return self.head(gate*l+(1-gate)*g)
 92
 93def train(model, hybrid, xs, ys, ss, tt):
 94    opt=torch.optim.Adam(model.parameters(),lr=0.02)
 95    for _ in range(80):
 96        for x,y in zip(xs,ys):
 97            pred=model(x,ss,tt) if hybrid else model(x)
 98            loss=((pred-y)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
 99    with torch.no_grad():
100        vals=[]
101        for x,y in zip(xs,ys): vals.append(float(((model(x,ss,tt) if hybrid else model(x))-y).pow(2).mean()))
102    return float(np.mean(vals))
103
104xs, ys, ss, tt = make_data()
105dense_model=DenseModel().to(device); hybrid_model=HybridModel().to(device)
106dense_loss=train(dense_model,False,xs,ys,ss,tt); hybrid_loss=train(hybrid_model,True,xs,ys,ss,tt)
107
108# Timing uses one representative graph and identical projection width.
109timing=[]
110for nn in sizes:
111    z=torch.randn(nn,d,device=device); a,b=graph(nn,degree)
112    W=[torch.randn(d,width,device=device) for _ in range(3)]
113    td=timed(lambda: dense_attention(z,*W), repeats=5)
114    th=timed(lambda: (local_attention(z,a,b,*W), global_linear_attention(z,*W)), repeats=5)
115    timing.append({'N':nn,'dense_ms':td,'hybrid_ms':th,'speedup':td/th})
116
117result={'device':device,'predictions':{'local_mask_max_error':mask_error,'linear_reassociation_max_error':linear_error,'proxy_ratio_N64_to_N512':[ratios[0],ratios[-1]],'proxy_loglog_slope':slope},'learning_mse':{'dense':dense_loss,'hybrid':hybrid_loss},'timing':timing}
118print(json.dumps(result,indent=2))
119with open('results.json','w') as f: json.dump(result,f,indent=2)