1import json, math, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5import torch.nn.functional as F
  6
  7
  8def seed_all(s=7):
  9    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 10    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 11
 12
 13def get_device():
 14    if torch.cuda.is_available():
 15        try:
 16            torch.empty(1, device='cuda')
 17            return torch.device('cuda')
 18        except Exception:
 19            pass
 20    return torch.device('cpu')
 21
 22
 23class Profile(nn.Module):
 24    def __init__(self, hidden=16):
 25        super().__init__()
 26        self.net = nn.Sequential(nn.Linear(2, hidden), nn.Tanh(), nn.Linear(hidden, 1))
 27        nn.init.zeros_(self.net[-1].weight)
 28        nn.init.constant_(self.net[-1].bias, 0.0)
 29
 30    def forward(self, s, q):
 31        z = self.net(torch.stack((s, q), -1)).squeeze(-1)
 32        return F.softplus(z + 2.0) + 1e-4
 33
 34
 35class MetricAttention(nn.Module):
 36    def __init__(self, d, constrained=True):
 37        super().__init__()
 38        self.log_h = nn.Parameter(torch.zeros(d))
 39        self.b = nn.Parameter(torch.randn(d) * .08)
 40        self.profile = Profile()
 41        self.constrained = constrained
 42
 43    def metric(self, y, need_derivatives=False):
 44        H = F.softplus(self.log_h) + 1e-4
 45        alpha = torch.sqrt((H * y * y).sum(-1) + 1e-8)
 46        beta = (self.b * y).sum(-1)
 47        s = beta / (alpha + 1e-8)
 48        q = (self.b * self.b / H).sum()
 49        phi = self.profile(s, q.expand_as(s))
 50        if not need_derivatives:
 51            return alpha * phi, None
 52        # derivatives of scalar profile with respect to s
 53        ss = s.detach().requires_grad_(True)
 54        qq = q.detach().expand_as(ss)
 55        pp = self.profile(ss, qq)
 56        p1 = torch.autograd.grad(pp.sum(), ss, create_graph=True)[0]
 57        p2 = torch.autograd.grad(p1.sum(), ss, create_graph=True)[0]
 58        g = pp - ss*p1 + (q.detach() - ss*ss)*p2
 59        return alpha * phi, (g, phi, s, q, H)
 60
 61    def forward(self, Q, K, V, return_aux=False):
 62        y = Q[:, :, None, :] - K[:, None, :, :]
 63        cost, aux = self.metric(y, need_derivatives=return_aux)
 64        logits = -(cost * cost) / math.sqrt(Q.shape[-1])
 65        a = F.softmax(logits, -1)
 66        out = a @ V
 67        return out, a, aux
 68
 69
 70class DotAttention(nn.Module):
 71    def __init__(self, d):
 72        super().__init__(); self.scale = d ** -0.5
 73    def forward(self, Q, K, V, return_aux=False):
 74        a = F.softmax(Q @ K.transpose(-1, -2) * self.scale, -1)
 75        out = a @ V
 76        return out, a, None
 77
 78
 79class TinyModel(nn.Module):
 80    def __init__(self, vocab=32, d=24, constrained=True, kind='metric'):
 81        super().__init__(); self.emb=nn.Embedding(vocab,d); self.attn=DotAttention(d) if kind=='dot' else MetricAttention(d,constrained)
 82        self.ff=nn.Sequential(nn.Linear(d,d*2),nn.GELU(),nn.Linear(d*2,d)); self.norm=nn.LayerNorm(d); self.head=nn.Linear(d,vocab)
 83    def forward(self,x, return_aux=False):
 84        z=self.emb(x); o,a,aux=self.attn(z,z,z,return_aux); z=self.norm(z+o); z=self.norm(z+self.ff(z)); logits=self.head(z)
 85        return logits, a, aux
 86
 87
 88def convexity_check(dev):
 89    # Directly verify g for phi=1+c*s: g=1 and inspect numerical Hessian of F^2/2.
 90    d=3; c=.25; b=torch.tensor([.35,-.2,.1],device=dev); H=torch.tensor([1.2,.8,1.5],device=dev)
 91    vals=[]; mine=[]
 92    for _ in range(30):
 93        y=torch.randn(d,device=dev); y.requires_grad_()
 94        alpha=torch.sqrt((H*y*y).sum()); s=(b*y).sum()/alpha; q=(b*b/H).sum()
 95        phi=1+c*s; g=phi-s*c+(q-s*s)*0.0; vals.append(float(g))
 96        f=.5*(alpha*phi)**2
 97        grad=torch.autograd.grad(f,y,create_graph=True)[0]
 98        rows=[torch.autograd.grad(grad[i],y,retain_graph=True)[0] for i in range(d)]
 99        mine.append(float(torch.linalg.eigvalsh(torch.stack(rows)).min()))
100    # Compare with an intentionally nonconvex quadratic profile.
101    bad=[]
102    for _ in range(20):
103        y=torch.randn(d,device=dev); y.requires_grad_(); alpha=torch.sqrt((H*y*y).sum()); s=(b*y).sum()/alpha
104        f=.5*(alpha*(1-20*s*s))**2; gr=torch.autograd.grad(f,y,create_graph=True)[0]
105        hs=[torch.autograd.grad(gr[i],y,retain_graph=True)[0] for i in range(d)]
106        bad.append(float(torch.linalg.eigvalsh(torch.stack(hs)).min()))
107    return {'linear_profile_g_min':min(vals),'linear_profile_g_max':max(vals),'linear_profile_hessian_min':min(mine),'bad_profile_hessian_min':min(bad)}
108
109
110def run(kind, constrained, dev, seed=7, steps=180):
111    seed_all(seed); vocab=32; L=12; B=64
112    # deterministic modular next-token task with local context
113    base=torch.arange(L,device=dev)[None,:].repeat(B,1)
114    x=(base + torch.randint(0,vocab,(B,1),device=dev)) % vocab
115    y=(x+1)%vocab
116    model=TinyModel(vocab=vocab,kind=kind,constrained=constrained).to(dev)
117    opt=torch.optim.AdamW(model.parameters(),lr=3e-3)
118    losses=[]; neg=[]; ent=[]; spikes=[]
119    for step in range(steps):
120        # refresh offsets, keeping the same simple task
121        x=(base + torch.randint(0,vocab,(B,1),device=dev)) % vocab; y=(x+1)%vocab
122        logits,a,aux=model(x,return_aux=True)
123        loss=F.cross_entropy(logits.reshape(-1,vocab),y.reshape(-1))
124        barrier=torch.tensor(0.,device=dev); frac=0.
125        if aux is not None and constrained:
126            g=aux[0]; barrier=F.softplus(.05-g).square().mean(); loss=loss+.03*barrier; frac=float((g<0).float().mean())
127        opt.zero_grad(); loss.backward(); gn=float(torch.nn.utils.clip_grad_norm_(model.parameters(),10.0)); opt.step()
128        losses.append(float(loss.detach())); neg.append(frac); spikes.append(gn)
129        p=a.detach().clamp_min(1e-8); ent.append(float((-(p*p.log()).sum(-1).mean()).cpu()))
130    with torch.no_grad():
131        logits,_,aux=model(x,return_aux=False); val=float(F.cross_entropy(logits.reshape(-1,vocab),y.reshape(-1)))
132        if aux is not None: neg_final=float((aux[0]<0).float().mean())
133        else: neg_final=0.
134    return {'final_loss':val,'mean_last20':float(np.mean(losses[-20:])),'loss_std_last50':float(np.std(losses[-50:])),'max_grad':max(spikes),'final_negative_g_fraction':neg_final,'mean_attention_entropy':float(np.mean(ent[-20:]))}
135
136
137def main():
138    seed_all(7); dev=get_device(); checks=convexity_check(dev)
139    results={'device':str(dev),'math_check':checks}
140    results['baseline_dot']=run('dot',False,dev,7)
141    results['unconstrained_metric']=run('metric',False,dev,7)
142    results['constrained_metric']=run('metric',True,dev,7)
143    with open('results.json','w') as f: json.dump(results,f,indent=2)
144    print(json.dumps(results,indent=2))
145
146if __name__=='__main__': main()