Upload exp5_mechanism.py with huggingface_hub

exp5_mechanism.py

@@ -0,0 +1,581 @@
+#!/usr/bin/env python3
+"""
+Experiment 5: Relaxation Mechanism Ablation
+
+Falsifies or confirms whether weight relaxation is structural prediction (BVP)
+or generic regularization. Five seeds throughout.
+
+Configs (all d=1024, 4L, 1000 steps, 10% active):
+  1. dense
+  2. dense + relax_graph (alpha=0.1)
+  3. dense + relax_roll (alpha=0.1)
+  4. ema_only
+  5. ema + relax_graph (alpha=0.1)
+  6. ema + relax_roll (alpha=0.1)
+  7. ema + relax_random (random similarity matrix)
+  8. ema + relax_shuffled_graph (real stats, broken structure)
+
+Alpha sweep (graph, ema): 0.0, 0.01, 0.05, 0.1, 0.2, 0.5, 0.9
+
+Diagnostics per run:
+  - Val loss every 50 steps
+  - grad_cos: cosine similarity between relaxer delta and dense oracle gradient
+  - mag_ratio: ||relaxer_delta|| / ||oracle_grad|| on inactive chunks
+  - mask_jaccard: step-to-step overlap of active set
+"""
+import argparse,json,math,os,random,sys,time,urllib.request
+from collections import defaultdict
+import torch,torch.nn as nn,torch.nn.functional as F
+import tiktoken
+print("imports ok",flush=True)
+
+# ═══════════════════════════════════════════════════════════════
+# DATA
+# ═══════════════════════════════════════════════════════════════
+class Corpus:
+    _i=None
+    @classmethod
+    def get(cls,bs,dev):
+        if cls._i is None: cls._i=cls(bs,dev)
+        return cls._i
+    def __init__(self,bs,dev):
+        self.block_size,self.device=bs,dev
+        p="input.txt"
+        if not os.path.exists(p):
+            urllib.request.urlretrieve("https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt",p)
+        enc=tiktoken.get_encoding("gpt2"); t=enc.encode(open(p).read())
+        self.vocab_size=enc.n_vocab; d=torch.tensor(t,dtype=torch.long)
+        si=int(0.9*len(d)); self.train_data,self.val_data=d[:si],d[si:]
+        print(f"Corpus: V={self.vocab_size} train={len(self.train_data):,} val={len(self.val_data):,}",flush=True)
+    def get_batch(self,split,bs,gen=None):
+        d=self.train_data if split=="train" else self.val_data
+        ix=torch.randint(len(d)-self.block_size-1,(bs,),generator=gen)
+        x=torch.stack([d[i:i+self.block_size] for i in ix])
+        y=torch.stack([d[i+1:i+self.block_size+1] for i in ix])
+        return x.to(self.device),y.to(self.device)
+def mg(s):
+    g=torch.Generator(device="cpu"); g.manual_seed(s); return g
+
+# ═══════════════════════════════════════════════════════════════
+# SPARSE LINEAR
+# ═══════════════════════════════════════════════════════════════
+class SparseBwd(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx,x,w,b,ac,cs,sdx):
+        ctx.save_for_backward(x,w,ac); ctx.hb=b is not None; ctx.sdx=sdx; ctx.cs=cs
+        return F.linear(x,w,b)
+    @staticmethod
+    def backward(ctx,gy):
+        x,w,ac=ctx.saved_tensors; cs=ctx.cs
+        xf=x.reshape(-1,x.shape[-1]); gf=gy.reshape(-1,gy.shape[-1])
+        gw=torch.zeros_like(w)
+        gb=torch.zeros(w.shape[0],device=w.device,dtype=w.dtype) if ctx.hb else None
+        gx=torch.zeros_like(xf) if ctx.sdx else gf@w
+        for c in ac.tolist():
+            s,e=c*cs,(c+1)*cs; sl=gf[:,s:e]
+            gw[s:e]=sl.t()@xf
+            if gb is not None: gb[s:e]=sl.sum(0)
+            if ctx.sdx: gx+=sl@w[s:e]
+        return gx.reshape(x.shape),gw,gb,None,None,None
+
+class SL(nn.Linear):
+    def __init__(self,i,o,bias=True):
+        super().__init__(i,o,bias=bias)
+        self.se=False; self.sdx=False; self.ac=None; self.cs=64
+    def forward(self,x):
+        if not self.se or self.ac is None: return F.linear(x,self.weight,self.bias)
+        return SparseBwd.apply(x,self.weight,self.bias,self.ac,self.cs,self.sdx)
+
+# ═══════════════════════════════════════════════════════════════
+# MODEL
+# ═══════════════════════════════════════════════════════════════
+class Attn(nn.Module):
+    def __init__(self,d,nh,bs,do):
+        super().__init__(); self.nh=nh; self.hd=d//nh
+        self.qkv=SL(d,3*d); self.proj=SL(d,d); self.drop=nn.Dropout(do)
+        self.register_buffer("mask",torch.tril(torch.ones(bs,bs)).view(1,1,bs,bs))
+    def forward(self,x):
+        B,T,C=x.shape; q,k,v=self.qkv(x).split(C,2)
+        q=q.view(B,T,self.nh,self.hd).transpose(1,2)
+        k=k.view(B,T,self.nh,self.hd).transpose(1,2)
+        v=v.view(B,T,self.nh,self.hd).transpose(1,2)
+        a=(q@k.transpose(-2,-1))/math.sqrt(self.hd)
+        a=a.masked_fill(self.mask[:,:,:T,:T]==0,float("-inf"))
+        a=self.drop(F.softmax(a,dim=-1))
+        return self.proj((a@v).transpose(1,2).contiguous().view(B,T,C))
+
+class FFN(nn.Module):
+    def __init__(self,d,do):
+        super().__init__(); self.fc=SL(d,4*d); self.proj=SL(4*d,d); self.drop=nn.Dropout(do)
+    def forward(self,x): return self.drop(self.proj(F.gelu(self.fc(x))))
+
+class Blk(nn.Module):
+    def __init__(self,d,nh,bs,do):
+        super().__init__(); self.ln1=nn.LayerNorm(d); self.attn=Attn(d,nh,bs,do)
+        self.ln2=nn.LayerNorm(d); self.mlp=FFN(d,do)
+    def forward(self,x): x=x+self.attn(self.ln1(x)); return x+self.mlp(self.ln2(x))
+
+class GPT(nn.Module):
+    def __init__(self,V,bs,nl,nh,d,do):
+        super().__init__(); self.te=nn.Embedding(V,d); self.pe=nn.Embedding(bs,d)
+        self.blocks=nn.Sequential(*[Blk(d,nh,bs,do) for _ in range(nl)])
+        self.ln=nn.LayerNorm(d); self.head=nn.Linear(d,V)
+    def forward(self,idx,tgt=None):
+        B,T=idx.shape; x=self.te(idx)+self.pe(torch.arange(T,device=idx.device))[None]
+        lo=self.head(self.ln(self.blocks(x)))
+        return lo,F.cross_entropy(lo.view(-1,lo.size(-1)),tgt.view(-1)) if tgt is not None else None
+    def np(self): return sum(p.numel() for p in self.parameters())
+def gsl(m): return [x for x in m.modules() if isinstance(x,SL)]
+
+# ═══════════════════════════════════════════════════════════════
+# SCHEDULER (builds similarity matrix during warmup)
+# ═══════════════════════════════════════════════════════════════
+class Sched:
+    def __init__(self,model,frac,cs,dev,beta=0.95,sim_hist=128,min_sim=8):
+        self.frac,self.cs,self.dev,self.beta=frac,cs,dev,beta
+        self.sim_hist,self.min_sim=sim_hist,min_sim
+        self.lins=gsl(model); self.m2i,self.m2l={},{}; off=0
+        for m in self.lins:
+            m.cs=cs; nc=m.out_features//cs; assert m.out_features%cs==0
+            self.m2i[m]=torch.arange(off,off+nc,device=dev)
+            self.m2l[m]=torch.arange(nc,device=dev); off+=nc
+        self.nc=off; self.ema=torch.zeros(self.nc,device=dev)
+        self.act=torch.zeros(self.nc,dtype=torch.bool,device=dev)
+        self.prev_act=torch.zeros(self.nc,dtype=torch.bool,device=dev)
+        self.mass_history=[]; self.similarity=None
+    def gf(self,step,wu,an):
+        if step<wu: return 1.0
+        if an>0 and step<wu+an:
+            p=(step-wu)/an; return self.frac+(1-self.frac)*0.5*(1+math.cos(math.pi*p))
+        return self.frac
+    def choose(self,step,wu,an):
+        self.prev_act=self.act.clone()
+        f=self.gf(step,wu,an)
+        if f>=0.999: self.act.fill_(True); self._inst(); return
+        k=max(1,int(f*self.nc)); self.act.fill_(False)
+        idx=torch.topk(self.ema+1e-9*torch.rand_like(self.ema),k=k).indices
+        self.act[idx]=True; self._inst()
+    def _inst(self):
+        for m,gi in self.m2i.items(): m.ac=self.m2l[m][self.act[gi]]
+    @torch.no_grad()
+    def update(self,step,wu):
+        cur=torch.zeros_like(self.ema)
+        for m,ids in self.m2i.items():
+            if m.weight.grad is None: continue
+            s=m.weight.grad.square().view(len(ids),self.cs,-1).sum((1,2))
+            if m.bias is not None and m.bias.grad is not None:
+                s+=m.bias.grad.square().view(len(ids),self.cs).sum(1)
+            cur[ids]=torch.sqrt(s+1e-30)
+        obs=self.act; new=obs&(self.ema==0); old=obs&~new
+        self.ema[new]=cur[new]; self.ema[old]=self.beta*self.ema[old]+(1-self.beta)*cur[old]
+        if step<wu:
+            self.mass_history.append(cur.clone())
+            if len(self.mass_history)>self.sim_hist:
+                self.mass_history=self.mass_history[-self.sim_hist:]
+            if len(self.mass_history)>=self.min_sim:
+                self._build_sim()
+        return cur
+    def _build_sim(self):
+        H=torch.stack(self.mass_history)
+        H=(H-H.mean(0,keepdim=True))/(H.std(0,keepdim=True)+1e-6)
+        S=torch.clamp((H.T@H)/max(1,H.shape[0]-1),min=0)
+        S.fill_diagonal_(0)
+        ok=torch.zeros_like(S,dtype=torch.bool)
+        for _,ids in self.m2i.items(): ok[ids[:,None],ids[None,:]]=True
+        self.similarity=torch.where(ok,S,torch.zeros_like(S))
+    def mask_jaccard(self):
+        """Jaccard between current and previous active set."""
+        if self.prev_act.sum()==0: return 0.0
+        i=(self.act&self.prev_act).sum().item()
+        u=(self.act|self.prev_act).sum().item()
+        return i/max(u,1)
+
+# ═══════════════════════════════════════════════════════════════
+# RELAXERS
+# ════════════════��══════════════════════════════════════════════
+class GraphRelaxer:
+    """Graph Laplacian relaxation using the real similarity matrix."""
+    def __init__(self, sched, alpha=0.1, iters=3):
+        self.sched,self.alpha,self.iters=sched,alpha,iters
+    @torch.no_grad()
+    def relax(self):
+        S=self.sched.similarity
+        if S is None: return {}
+        act=self.sched.act; deltas={}
+        for m,ids in self.sched.m2i.items():
+            nc=len(ids); cs=self.sched.cs; di=m.weight.shape[1]
+            S_local=S[ids][:,ids]
+            rs=S_local.sum(1,keepdim=True)+1e-12; S_n=S_local/rs
+            la=act[ids]; li=~la
+            if li.sum()==0: continue
+            W=m.weight.data.view(nc,cs,di); W_before=W[li].clone()
+            for _ in range(self.iters):
+                Wf=W.reshape(nc,-1); Wa=(S_n@Wf).view(nc,cs,di)
+                W[li]=(1-self.alpha)*W[li]+self.alpha*Wa[li]
+            m.weight.data=W.view(m.out_features,di)
+            deltas[m]=W[li]-W_before  # (n_inactive, cs, di)
+        return deltas
+
+class RollRelaxer:
+    """Spatial neighbor relaxation via torch.roll."""
+    def __init__(self, sched, alpha=0.1, iters=3):
+        self.sched,self.alpha,self.iters=sched,alpha,iters
+    @torch.no_grad()
+    def relax(self):
+        act=self.sched.act; deltas={}
+        for m,ids in self.sched.m2i.items():
+            nc=len(ids); cs=self.sched.cs; di=m.weight.shape[1]
+            la=act[ids]; li=~la
+            if li.sum()==0: continue
+            W=m.weight.data.view(nc,cs,di); W_before=W[li].clone()
+            for _ in range(self.iters):
+                Wp=torch.roll(W,1,dims=0); Wn=torch.roll(W,-1,dims=0)
+                Wa=(Wp+Wn)/2.0
+                W[li]=(1-self.alpha)*W[li]+self.alpha*Wa[li]
+            m.weight.data=W.view(m.out_features,di)
+            deltas[m]=W[li]-W_before
+        return deltas
+
+class RandomRelaxer:
+    """Control: random similarity matrix (same sparsity pattern, random values)."""
+    def __init__(self, sched, alpha=0.1, iters=3):
+        self.sched,self.alpha,self.iters=sched,alpha,iters
+        self._rand_sim=None
+    def _get_rand_sim(self):
+        if self._rand_sim is not None: return self._rand_sim
+        S=self.sched.similarity
+        if S is None: return None
+        # Random positive values with same mask structure
+        R=torch.rand_like(S)*S.abs().mean()
+        R.fill_diagonal_(0)
+        ok=torch.zeros_like(R,dtype=torch.bool)
+        for _,ids in self.sched.m2i.items(): ok[ids[:,None],ids[None,:]]=True
+        self._rand_sim=torch.where(ok,R,torch.zeros_like(R))
+        return self._rand_sim
+    @torch.no_grad()
+    def relax(self):
+        S=self._get_rand_sim()
+        if S is None: return {}
+        act=self.sched.act; deltas={}
+        for m,ids in self.sched.m2i.items():
+            nc=len(ids); cs=self.sched.cs; di=m.weight.shape[1]
+            S_local=S[ids][:,ids]
+            rs=S_local.sum(1,keepdim=True)+1e-12; S_n=S_local/rs
+            la=act[ids]; li=~la
+            if li.sum()==0: continue
+            W=m.weight.data.view(nc,cs,di); W_before=W[li].clone()
+            for _ in range(self.iters):
+                Wf=W.reshape(nc,-1); Wa=(S_n@Wf).view(nc,cs,di)
+                W[li]=(1-self.alpha)*W[li]+self.alpha*Wa[li]
+            m.weight.data=W.view(m.out_features,di)
+            deltas[m]=W[li]-W_before
+        return deltas
+
+class ShuffledGraphRelaxer:
+    """Control: real similarity stats, shuffled structure within each layer."""
+    def __init__(self, sched, alpha=0.1, iters=3):
+        self.sched,self.alpha,self.iters=sched,alpha,iters
+        self._shuf_sim=None
+    def _get_shuf_sim(self):
+        if self._shuf_sim is not None: return self._shuf_sim
+        S=self.sched.similarity
+        if S is None: return None
+        Ss=S.clone()
+        # Shuffle within each layer block
+        for _,ids in self.sched.m2i.items():
+            n=len(ids)
+            block=Ss[ids][:,ids].clone()  # (n,n)
+            # Shuffle rows and columns with same permutation
+            perm=torch.randperm(n,device=S.device)
+            block=block[perm][:,perm]
+            block.fill_diagonal_(0)
+            Ss[ids[:,None],ids[None,:]]=block
+        self._shuf_sim=Ss
+        return self._shuf_sim
+    @torch.no_grad()
+    def relax(self):
+        S=self._get_shuf_sim()
+        if S is None: return {}
+        act=self.sched.act; deltas={}
+        for m,ids in self.sched.m2i.items():
+            nc=len(ids); cs=self.sched.cs; di=m.weight.shape[1]
+            S_local=S[ids][:,ids]
+            rs=S_local.sum(1,keepdim=True)+1e-12; S_n=S_local/rs
+            la=act[ids]; li=~la
+            if li.sum()==0: continue
+            W=m.weight.data.view(nc,cs,di); W_before=W[li].clone()
+            for _ in range(self.iters):
+                Wf=W.reshape(nc,-1); Wa=(S_n@Wf).view(nc,cs,di)
+                W[li]=(1-self.alpha)*W[li]+self.alpha*Wa[li]
+            m.weight.data=W.view(m.out_features,di)
+            deltas[m]=W[li]-W_before
+        return deltas
+
+class NullRelaxer:
+    """No-op relaxer."""
+    def relax(self): return {}
+
+# ═══════════════════════════════════════════════════════════════
+# OPTIMIZER
+# ═══════════════════════════════════════════════════════════════
+class CAdam:
+    def __init__(self,model,lr=3e-4,cs=64):
+        self.model,self.lr,self.cs=model,lr,cs
+        self.st={}; self.p2m={}
+        for m in gsl(model):
+            if m.weight is not None: self.p2m[m.weight]=m
+            if m.bias is not None: self.p2m[m.bias]=m
+    def zero_grad(self):
+        for p in self.model.parameters(): p.grad=None
+    @torch.no_grad()
+    def step(self):
+        for p in self.model.parameters():
+            if p.grad is None: continue
+            if p not in self.st: self.st[p]={"m":torch.zeros_like(p),"v":torch.zeros_like(p)}
+            m,v=self.st[p]["m"],self.st[p]["v"]
+            sm=self.p2m.get(p); ac=getattr(sm,'ac',None) if sm else None
+            if ac is None:
+                m.mul_(0.9).add_(p.grad,alpha=0.1); v.mul_(0.999).addcmul_(p.grad,p.grad,value=0.001)
+                p.sub_(m/(torch.sqrt(v)+1e-8),alpha=self.lr)
+            else:
+                m.mul_(0.9).add_(p.grad,alpha=0.1); v.mul_(0.999).addcmul_(p.grad,p.grad,value=0.001)
+                for c in ac.tolist():
+                    s,e=c*self.cs,(c+1)*self.cs
+                    p.data[s:e].sub_(m[s:e]/(torch.sqrt(v[s:e])+1e-8),alpha=self.lr)
+
+# ═══════════════════════════════════════════════════════════════
+# EVAL
+# ═══════════════════════════════════════════════════════════════
+@torch.no_grad()
+def ev(model,corpus,bs,n=20,seed=9999):
+    model.eval(); ls=[model(*corpus.get_batch("val",bs,mg(seed+i)))[1].item() for i in range(n)]
+    model.train(); a=sum(ls)/len(ls); return a,math.exp(min(a,20))
+
+# ═══════════════════════════════════════════════════════════════
+# ORACLE GRADIENT DIAGNOSTIC
+# ═══════════════════════════════════════════════════════════════
+@torch.no_grad()
+def compute_relaxer_diagnostics(model, sched, relaxer_deltas, x, y, corpus, bs, cs):
+    """
+    Compare relaxer delta on inactive chunks to what dense gradient would have been.
+    Returns (grad_cos, mag_ratio) or (None, None) if not applicable.
+    """
+    if not relaxer_deltas: return None, None
+
+    # Compute dense gradients
+    for m in gsl(model): m.se=False
+    for p in model.parameters(): p.grad=None
+    _,lo=model(x,y); lo.backward()
+
+    cos_sims=[]; mag_ratios=[]
+    for m,delta in relaxer_deltas.items():
+        if m not in sched.m2i: continue
+        ids=sched.m2i[m]; nc=len(ids); di=m.weight.shape[1]
+        la=sched.act[ids]; li=~la
+        if li.sum()==0 or m.weight.grad is None: continue
+
+        # Dense gradient for inactive chunks, reshaped
+        dense_g=m.weight.grad.view(nc,cs,di)[li]  # (n_inact, cs, di)
+
+        # Flatten for cosine/magnitude
+        d_flat=delta.reshape(-1); g_flat=dense_g.reshape(-1)
+        dn=d_flat.norm(); gn=g_flat.norm()
+        if dn>1e-12 and gn>1e-12:
+            cos_sims.append(F.cosine_similarity(d_flat.unsqueeze(0),g_flat.unsqueeze(0)).item())
+            mag_ratios.append((dn/gn).item())
+
+    # Restore sparse mode
+    for m in gsl(model): m.se=True
+    for p in model.parameters(): p.grad=None
+
+    if not cos_sims: return None, None
+    return sum(cos_sims)/len(cos_sims), sum(mag_ratios)/len(mag_ratios)
+
+# ═══════════════════════════════════════════════════════════════
+# SINGLE RUN
+# ═══════════════════════════════════════════════════════════════
+def run1(mode, steps, bs, bsz, nl, nh, d, cs, af, wu, an, lr, dev, seed,
+         alpha=0.1, iters=3, diag_interval=100):
+    torch.manual_seed(seed); random.seed(seed)
+    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
+    corpus=Corpus.get(bsz,dev)
+    model=GPT(corpus.vocab_size,bsz,nl,nh,d,0.1).to(dev)
+    for m in gsl(model): m.cs=cs
+
+    is_dense = mode.startswith("dense")
+    is_sparse = not is_dense
+    needs_relax = "relax" in mode
+
+    sched=None
+    if is_sparse:
+        sched=Sched(model,af,cs,dev)
+    elif needs_relax:
+        # Dense + relax: need scheduler for similarity matrix but run dense forward/backward
+        sched=Sched(model,af,cs,dev)
+
+    opt=CAdam(model,lr,cs)
+
+    # Create relaxer
+    if not needs_relax:
+        relaxer=NullRelaxer()
+    elif "random" in mode:
+        relaxer=RandomRelaxer(sched,alpha,iters)
+    elif "shuffled" in mode:
+        relaxer=ShuffledGraphRelaxer(sched,alpha,iters)
+    elif "roll" in mode:
+        relaxer=RollRelaxer(sched,alpha,iters)
+    elif "graph" in mode:
+        relaxer=GraphRelaxer(sched,alpha,iters)
+    else:
+        relaxer=NullRelaxer()
+
+    np_=model.np()
+    val_curve=[]; grad_cos_log=[]; mag_ratio_log=[]; jaccard_log=[]
+
+    if dev=="cuda": torch.cuda.synchronize()
+    t0=time.perf_counter()
+
+    for step in range(steps):
+        x,y=corpus.get_batch("train",bs,mg(step))
+
+        if is_sparse:
+            sched.choose(step,wu,an)
+            for m in gsl(model): m.se=True; m.sdx=False
+        else:
+            for m in gsl(model): m.se=False; m.ac=None
+            # For dense+relax: still run scheduler to build similarity & set active mask
+            if sched:
+                sched.choose(step,wu,an)
+
+        opt.zero_grad(); _,loss=model(x,y); loss.backward()
+
+        if sched:
+            sched.update(step,wu)
+            jaccard_log.append((step, sched.mask_jaccard()))
+
+        opt.step()
+
+        # Relaxation (only after annealing completes)
+        relax_deltas={}
+        if needs_relax and step>=wu+an:
+            # For dense+relax: temporarily set active mask so relaxer knows what's "active"
+            if is_dense and sched:
+                for m,ids in sched.m2i.items():
+                    m.ac=sched.m2l[m][sched.act[ids]]
+            relax_deltas=relaxer.relax()
+            if is_dense and sched:
+                for m in gsl(model): m.ac=None
+
+        # Diagnostics
+        if step%50==0:
+            vl,_=ev(model,corpus,bs,n=10,seed=7777)
+            val_curve.append((step,vl))
+
+        if step%diag_interval==0 and step>=wu+an and relax_deltas and sched:
+            gc,mr=compute_relaxer_diagnostics(model,sched,relax_deltas,x,y,corpus,bs,cs)
+            if gc is not None:
+                grad_cos_log.append((step,gc))
+                mag_ratio_log.append((step,mr))
+
+        if step%200==0:
+            print(f"    step {step}/{steps} loss={loss.item():.4f}",flush=True)
+
+    if dev=="cuda": torch.cuda.synchronize()
+    wall=time.perf_counter()-t0
+    for m in gsl(model): m.se=False
+    vl,vp=ev(model,corpus,bs,n=30)
+    del model; torch.cuda.empty_cache() if dev=="cuda" else None
+
+    return {
+        "vl":vl,"vp":vp,"wall":wall,"ms":1000*wall/steps,"np":np_,"tl":loss.item(),
+        "val_curve":val_curve,"grad_cos":grad_cos_log,"mag_ratio":mag_ratio_log,
+        "jaccard":jaccard_log,
+    }
+
+def runs(cfg,seeds):
+    rs=[]
+    for s in seeds:
+        c=dict(cfg); c["seed"]=s; rs.append(run1(**c))
+    vls=[r["vl"] for r in rs]; ml=sum(vls)/len(vls)
+    sl=(sum((x-ml)**2 for x in vls)/max(1,len(vls)-1))**0.5
+    return {"ml":ml,"sl":sl,"rs":rs,"ms":sum(r["ms"] for r in rs)/len(rs)}
+
+# ═══════════════════════════════════════════════════════════════
+# MAIN
+# ═══════════════════════════════════════════════════════════════
+def main():
+    p=argparse.ArgumentParser()
+    p.add_argument("--device",default="cuda"); p.add_argument("--steps",type=int,default=1000)
+    p.add_argument("--seeds",default="42,123,456,789,1024")
+    p.add_argument("--d",type=int,default=1024); p.add_argument("--nl",type=int,default=4)
+    p.add_argument("--nh",type=int,default=8); p.add_argument("--bs",type=int,default=8)
+    p.add_argument("--bsz",type=int,default=256); p.add_argument("--cs",type=int,default=64)
+    p.add_argument("--af",type=float,default=0.10); p.add_argument("--wu",type=int,default=50)
+    p.add_argument("--an",type=int,default=200); p.add_argument("--lr",type=float,default=3e-4)
+    a=p.parse_args(); seeds=[int(s) for s in a.seeds.split(",")]
+
+    if a.device=="cuda" and torch.cuda.is_available():
+        print(f"GPU: {torch.cuda.get_device_name()} VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB",flush=True)
+    print(f"d={a.d} nl={a.nl} steps={a.steps} seeds={seeds}",flush=True)
+    print(f"cs={a.cs} af={a.af} wu={a.wu} an={a.an} lr={a.lr}",flush=True)
+
+    base=dict(steps=a.steps,bs=a.bs,bsz=a.bsz,nl=a.nl,nh=a.nh,d=a.d,cs=a.cs,af=a.af,
+              wu=a.wu,an=a.an,lr=a.lr,dev=a.device,alpha=0.1,iters=3)
+
+    # ── Part 1: Main configs ──
+    configs=[
+        ("dense",                "dense"),
+        ("dense+relax_graph",    "dense+relax_graph"),
+        ("dense+relax_roll",     "dense+relax_roll"),
+        ("ema_only",             "ema_only"),
+        ("ema+relax_graph",      "ema+relax_graph"),
+        ("ema+relax_roll",       "ema+relax_roll"),
+        ("ema+relax_random",     "ema+relax_random"),
+        ("ema+relax_shuffled",   "ema+relax_shuffled_graph"),
+    ]
+
+    print("\n"+"="*80,flush=True)
+    print("EXP 5: Relaxation Mechanism Ablation (5 seeds)",flush=True)
+    print("="*80,flush=True)
+
+    R={}
+    for name,mode in configs:
+        print(f"\n--- {name} ({len(seeds)} seeds) ---",flush=True)
+        R[name]=runs({**base,"mode":mode},seeds)
+
+    print(f"\n{'Method':<25} | {'Val Loss':>20} | {'ms/step':>8}",flush=True)
+    print("-"*60,flush=True)
+    for name,_ in configs:
+        r=R[name]
+        print(f"{name:<25} | {r['ml']:.4f} ± {r['sl']:.4f}      | {r['ms']:>7.1f}",flush=True)
+
+    # ── Part 2: Alpha sweep ──
+    print(f"\n--- Alpha sweep (ema+relax_graph, 5 seeds) ---",flush=True)
+    print(f"{'alpha':>6} | {'Val Loss':>20} | {'ms/step':>8}",flush=True)
+    print("-"*42,flush=True)
+    alpha_results={}
+    for alpha in [0.0, 0.01, 0.05, 0.1, 0.2, 0.5, 0.9]:
+        r=runs({**base,"mode":"ema+relax_graph","alpha":alpha},seeds)
+        alpha_results[alpha]=r
+        print(f"{alpha:>6.2f} | {r['ml']:.4f} ± {r['sl']:.4f}      | {r['ms']:>7.1f}",flush=True)
+
+    # ── Part 3: Diagnostics summary ──
+    print(f"\n--- Diagnostics (grad_cos, mag_ratio) ---",flush=True)
+    for name in ["ema+relax_graph","ema+relax_roll","ema+relax_random","ema+relax_shuffled"]:
+        if name not in R: continue
+        gc_all=[]; mr_all=[]
+        for res in R[name]["rs"]:
+            gc_all.extend([x[1] for x in res["grad_cos"]])
+            mr_all.extend([x[1] for x in res["mag_ratio"]])
+        if gc_all:
+            gc_m=sum(gc_all)/len(gc_all); mr_m=sum(mr_all)/len(mr_all)
+            print(f"  {name:<25}: grad_cos={gc_m:.4f}  mag_ratio={mr_m:.4f}",flush=True)
+
+    # Save
+    all_results={"configs":R,"alpha_sweep":alpha_results}
+    with open("exp5.json","w") as f:
+        json.dump(all_results,f,indent=2,default=str)
+    print("\n✓ exp5.json saved",flush=True)
+    print(f"\nTotal: {len(configs)*len(seeds) + 7*len(seeds)} runs",flush=True)
+
+if __name__=="__main__": main()