Upload run_graduated.py with huggingface_hub

run_graduated.py

@@ -0,0 +1,509 @@
+#!/usr/bin/env python3
+"""
+Graduated Dissimilarity Experiment
+====================================
+Train on modular addition, then fork into 5 branches with tasks of
+increasing dissimilarity. Track representation metrics to find the
+tipping point where forgetting begins.
+
+Branches:
+  A→A: Continue addition        (Level 0: identical)
+  A→B: Switch to subtraction    (Level 1: same Fourier circuit)
+  A→C: Switch to multiplication (Level 2: different Fourier freqs)
+  A→D: Switch to max(a,b)       (Level 3: linear/ordinal circuit)
+  A→E: Switch to XOR            (Level 4: bit-level circuit)
+"""
+
+import torch, torch.nn as nn, torch.optim as optim, torch.nn.functional as F
+import numpy as np, json, os, math, copy, time
+from dataclasses import dataclass
+from torch.utils.data import Dataset, DataLoader
+from typing import Dict, List, Tuple
+
+# ============= INLINE DEPS (self-contained for job) =============
+
+# --- Representation metrics ---
+def centering(K):
+    n=K.shape[0]; u=torch.ones(n,n,device=K.device,dtype=K.dtype)/n
+    return K-u@K-K@u+u@K@u
+def linear_HSIC(X,Y):
+    K=X@X.T; L=Y@Y.T; return (centering(K)*centering(L)).sum()/((X.shape[0]-1)**2)
+def linear_CKA(X,Y):
+    xy=linear_HSIC(X,Y);xx=linear_HSIC(X,X);yy=linear_HSIC(Y,Y)
+    return (xy/(xx.sqrt()*yy.sqrt()).clamp(min=1e-10)).item()
+def cka_heatmap(sa,sb):
+    hm=np.zeros((len(sa),len(sb)))
+    for i in range(len(sa)):
+        for j in range(len(sb)): hm[i,j]=linear_CKA(sa[i],sb[j])
+    return hm
+def subspace_angles(X,Y,k=10):
+    def tb(Z,k):
+        _,_,Vh=torch.linalg.svd(Z-Z.mean(0),full_matrices=False); return Vh[:min(k,Vh.shape[0])].T
+    Qx=tb(X,k);Qy=tb(Y,k); mk=min(Qx.shape[1],Qy.shape[1])
+    return torch.arccos(torch.linalg.svdvals(Qx[:,:mk].T@Qy[:,:mk]).clamp(-1,1))
+def mean_sa_deg(X,Y,k=10): return (subspace_angles(X,Y,k).mean()*180/torch.pi).item()
+def grad_align(model,ba,bb,lfn):
+    model.zero_grad();lfn(model,ba).backward()
+    ga=torch.cat([p.grad.flatten() for p in model.parameters() if p.grad is not None]).clone()
+    model.zero_grad();lfn(model,bb).backward()
+    gb=torch.cat([p.grad.flatten() for p in model.parameters() if p.grad is not None]).clone()
+    model.zero_grad(); return F.cosine_similarity(ga.unsqueeze(0),gb.unsqueeze(0)).item()
+def attn_entropy(aw):
+    H=-(aw*(aw+1e-9).log2()).sum(-1)
+    return {'mean':H.mean().item(),'per_head':H.mean(dim=(0,2)).cpu().tolist()}
+def wc_mag(s0,s1):
+    return {n:(s1[n].float()-s0[n].float()).norm().item() for n in s0 if n in s1}
+
+# Fourier power spectrum of embedding matrix
+def fourier_power_spectrum(W_E, p):
+    """Compute Fourier power at each frequency for token embeddings.
+    W_E: [vocab_size, d_model]. Returns power at each freq 0..p//2."""
+    # Extract number embeddings only (skip special tokens)
+    from tasks import NUM_SPECIAL
+    num_emb = W_E[NUM_SPECIAL:NUM_SPECIAL+p, :]  # [p, d_model]
+    fft = torch.fft.rfft(num_emb, dim=0)  # [p//2+1, d_model]
+    power = (fft.abs() ** 2).sum(dim=1)  # [p//2+1]
+    return power.cpu().numpy()
+
+# --- Tasks ---
+NS = 7  # NUM_SPECIAL (updated)
+DP = 97
+
+OP_TOKENS = {'add':2,'subtract':3,'multiply':4,'max':5,'xor':6}
+ALL_OPS = ['add','subtract','multiply','max','xor']
+
+class MAD(Dataset):
+    def __init__(self,op='add',p=DP,split='train',tf=0.5,seed=42):
+        self.p=p;self.op=op;self.ot=OP_TOKENS[op]
+        ap=[(a,b) for a in range(p) for b in range(p)]
+        rng=np.random.RandomState(seed);rng.shuffle(ap)
+        nt=int(len(ap)*tf); self.pairs=ap[:nt] if split=='train' else ap[nt:]
+    def _c(self,a,b):
+        if self.op=='add': return(a+b)%self.p
+        elif self.op=='subtract': return(a-b)%self.p
+        elif self.op=='multiply': return(a*b)%self.p
+        elif self.op=='max': return max(a,b)
+        elif self.op=='xor': return(a^b)%self.p
+    def __len__(self): return len(self.pairs)
+    def __getitem__(self,i):
+        a,b=self.pairs[i];c=self._c(a,b)
+        return {'input_ids':torch.tensor([a+NS,self.ot,b+NS,1,c+NS],dtype=torch.long),
+                'labels':torch.tensor([-100,-100,-100,-100,c+NS],dtype=torch.long)}
+
+def get_probe(ds,n=500):
+    n=min(n,len(ds));its=[ds[i] for i in range(n)]
+    return torch.stack([it['input_ids'] for it in its]),np.array([it['labels'][-1].item()-NS for it in its])
+
+def get_loaders(p=DP,bs=512,tf=0.5,seed=42):
+    ld={}
+    for op in ALL_OPS:
+        for sp in ['train','test']:
+            ld[f'{op}_{sp}']=DataLoader(MAD(op,p,sp,tf,seed),batch_size=bs,shuffle=(sp=='train'),drop_last=False)
+    return ld
+
+# --- Model ---
+@dataclass
+class TC:
+    vs:int=104;nl:int=2;dm:int=128;nh:int=4;dmlp:int=512;msl:int=5;do:float=0.0;ln:bool=True
+
+class MHA(nn.Module):
+    def __init__(self,c):
+        super().__init__();self.nh=c.nh;self.dh=c.dm//c.nh
+        self.WQ=nn.Linear(c.dm,c.dm,bias=False);self.WK=nn.Linear(c.dm,c.dm,bias=False)
+        self.WV=nn.Linear(c.dm,c.dm,bias=False);self.WO=nn.Linear(c.dm,c.dm,bias=False)
+    def forward(self,x,ra=False):
+        B,T,D=x.shape
+        Q=self.WQ(x).view(B,T,self.nh,self.dh).transpose(1,2)
+        K=self.WK(x).view(B,T,self.nh,self.dh).transpose(1,2)
+        V=self.WV(x).view(B,T,self.nh,self.dh).transpose(1,2)
+        s=(Q@K.transpose(-2,-1))/math.sqrt(self.dh)
+        s.masked_fill_(torch.triu(torch.ones(T,T,device=x.device),1).bool().unsqueeze(0).unsqueeze(0),float('-inf'))
+        a=F.softmax(s,dim=-1);return self.WO((a@V).transpose(1,2).reshape(B,T,D)),(a if ra else None)
+
+class MLP2(nn.Module):
+    def __init__(self,c):
+        super().__init__();self.wi=nn.Linear(c.dm,c.dmlp);self.wo=nn.Linear(c.dmlp,c.dm)
+    def forward(self,x):h=F.gelu(self.wi(x));return self.wo(h),h
+
+class TB(nn.Module):
+    def __init__(self,c):
+        super().__init__();self.attn=MHA(c);self.mlp=MLP2(c)
+        self.ln1=nn.LayerNorm(c.dm) if c.ln else nn.Identity()
+        self.ln2=nn.LayerNorm(c.dm) if c.ln else nn.Identity()
+    def forward(self,x,ri=False):
+        ao,aw=self.attn(self.ln1(x),ra=ri);x=x+ao;mo,mh=self.mlp(self.ln2(x));x=x+mo
+        r={'hs':x}
+        if ri:r['aw']=aw;r['mh']=mh
+        return r
+
+class ST(nn.Module):
+    def __init__(self,c):
+        super().__init__();self.c=c
+        self.te=nn.Embedding(c.vs,c.dm);self.pe=nn.Embedding(c.msl,c.dm)
+        self.blocks=nn.ModuleList([TB(c) for _ in range(c.nl)])
+        self.lnf=nn.LayerNorm(c.dm) if c.ln else nn.Identity()
+        self.head=nn.Linear(c.dm,c.vs,bias=False);self.head.weight=self.te.weight
+        self.apply(self._iw)
+    def _iw(self,m):
+        if isinstance(m,nn.Linear):nn.init.normal_(m.weight,std=0.02)
+        elif isinstance(m,nn.Embedding):nn.init.normal_(m.weight,std=0.02)
+    def forward(self,ids,labels=None,ri=False):
+        B,T=ids.shape;x=self.te(ids)+self.pe(torch.arange(T,device=ids.device))
+        ahs=[x.detach()];aaw=[];amh=[]
+        for b in self.blocks:
+            r=b(x,ri=ri);x=r['hs'];ahs.append(x.detach())
+            if ri:aaw.append(r['aw'].detach());amh.append(r['mh'].detach())
+        lo=self.head(self.lnf(x));res={'logits':lo}
+        if labels is not None:res['loss']=F.cross_entropy(lo.view(-1,lo.size(-1)),labels.view(-1),ignore_index=-100)
+        if ri:res['hs']=ahs;res['aw']=aaw;res['mh']=amh
+        return res
+    def nparams(self):return sum(p.numel() for p in self.parameters())
+
+# --- Experiment helpers ---
+def ev(model,dl,dev):
+    model.eval();tl=tc=tt=0
+    with torch.no_grad():
+        for b in dl:
+            ids=b['input_ids'].to(dev);labs=b['labels'].to(dev)
+            o=model(ids,labels=labs);tl+=o['loss'].item()*ids.shape[0]
+            tc+=(o['logits'][:,-1,:].argmax(-1)==labs[:,-1]).sum().item();tt+=ids.shape[0]
+    return {'loss':tl/tt,'acc':tc/tt}
+
+def creps(model,pids,dev,pos=-1):
+    model.eval()
+    with torch.no_grad():o=model(pids.to(dev),ri=True)
+    return {'hs':[h[:,pos,:].cpu() for h in o['hs']],'aw':[a.cpu() for a in o['aw']]}
+
+def cmetrics(model,s0,s1,rc,ri,rp,dev,cfg):
+    m={};nl=cfg.nl+1
+    for li in range(nl):
+        p=f'layer_{li}';c=rc['hs'][li];ii=ri['hs'][li];pp=rp['hs'][li]
+        m[f'{p}/cka_vs_init']=linear_CKA(c,ii)
+        m[f'{p}/cka_vs_phase1']=linear_CKA(c,pp)
+        k=min(10,c.shape[0]//2,c.shape[1])
+        m[f'{p}/subspace_angle_vs_phase1']=mean_sa_deg(c,pp,k=k) if k>0 else 0.
+    for li,aw in enumerate(rc['aw']):
+        e=attn_entropy(aw);m[f'layer_{li+1}/attn_entropy_mean']=e['mean']
+        for h,he in enumerate(e['per_head']):m[f'layer_{li+1}/head_{h}_entropy']=he
+    cs={k:v.cpu() for k,v in model.state_dict().items()}
+    wp=wc_mag(s1,cs)
+    for bi in range(cfg.nl):
+        m[f'block_{bi}/weight_change_from_phase1']=sum(v for k,v in wp.items() if f'blocks.{bi}' in k)
+    return m
+
+def train_phase(model,opt,dl,ne,dev,pn,s0,s1,ri,rp,pids,eval_loaders,cfg,ce=20):
+    hist=[];gs=0
+    for ep in range(ne):
+        model.train();el=nb=0
+        for b in dl:
+            ids=b['input_ids'].to(dev);labs=b['labels'].to(dev)
+            o=model(ids,labels=labs);o['loss'].backward();opt.step();opt.zero_grad()
+            el+=o['loss'].item();nb+=1;gs+=1
+            if gs%ce==0:
+                model.eval();rc=creps(model,pids,dev)
+                sm=cmetrics(model,s0,s1,rc,ri,rp,dev,cfg)
+                # Evaluate on ALL tasks
+                for n,ld in eval_loaders.items():
+                    if '_test' in n:
+                        e=ev(model,ld,dev);sm[f'eval/{n}_loss']=e['loss'];sm[f'eval/{n}_acc']=e['acc']
+                # Gradient alignment: addition vs each other task
+                add_batch=next(iter(eval_loaders['add_test']))
+                def lfn(m,b):return m(b['input_ids'].to(dev),labels=b['labels'].to(dev))['loss']
+                for task_name in ['subtract','multiply','max','xor']:
+                    try:
+                        tb=next(iter(eval_loaders[f'{task_name}_test']))
+                        sm[f'grad_align_add_vs_{task_name}']=grad_align(model,add_batch,tb,lfn)
+                    except:sm[f'grad_align_add_vs_{task_name}']=0.
+                # Fourier power spectrum of embeddings
+                emb_w = model.te.weight.detach().cpu()
+                fps = fourier_power_spectrum(emb_w, cfg.vs - NS)
+                # Store top-5 peak frequencies
+                top5 = np.argsort(fps)[::-1][:5]
+                sm['fourier_top5_freqs'] = top5.tolist()
+                sm['fourier_top5_power'] = fps[top5].tolist()
+                sm['fourier_total_power'] = float(fps.sum())
+                sm['fourier_concentration'] = float(fps[top5].sum() / fps.sum())  # how concentrated
+
+                sm['phase']=pn;sm['epoch']=ep;sm['step']=gs;sm['train_loss']=el/nb
+                hist.append(sm)
+                add_acc=sm.get('eval/add_test_acc',0)
+                task_acc=sm.get(f'eval/{pn.split("_")[-1]}_test_acc', add_acc)
+                print(f"[{pn}] S{gs} L:{el/nb:.4f} AddAcc:{add_acc:.3f} TaskAcc:{task_acc:.3f} "
+                      f"CKA(L1vP1):{sm.get('layer_1/cka_vs_phase1',0):.3f} "
+                      f"FourierConc:{sm.get('fourier_concentration',0):.3f}")
+                model.train()
+        print(f"[{pn}] Ep{ep+1}/{ne} L:{el/nb:.4f}")
+    return hist
+
+# ============= MAIN =============
+print("="*70)
+print("GRADUATED DISSIMILARITY EXPERIMENT")
+print("="*70)
+t0=time.time()
+dev=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(f"Device: {dev}")
+
+p=97;seed=42;torch.manual_seed(seed);np.random.seed(seed)
+cfg=TC(vs=p+NS,nl=2,dm=128,nh=4,dmlp=512,msl=5)
+model=ST(cfg).to(dev)
+print(f"Parameters: {model.nparams():,}")
+
+s0={k:v.cpu().clone() for k,v in model.state_dict().items()}
+lds=get_loaders(p=p,bs=512,tf=0.5,seed=seed)
+
+# Probe data from addition test set (same across all branches)
+dsa=MAD('add',p=p,split='test',tf=0.5,seed=seed)
+pids_a,pla=get_probe(dsa,500)
+ri=creps(model,pids_a,dev)
+
+# Initial Fourier spectrum
+init_fps = fourier_power_spectrum(model.te.weight.detach().cpu(), p)
+
+# ========== PHASE 1: Train on Addition ==========
+print(f"\n{'='*60}\nPHASE 1: Train on Addition ({150} epochs)\n{'='*60}")
+opt=optim.AdamW(model.parameters(),lr=1e-3,weight_decay=1.0)
+h1=train_phase(model,opt,lds['add_train'],150,dev,'p1_add',s0,s0,ri,ri,pids_a,lds,cfg,ce=20)
+
+s1={k:v.cpu().clone() for k,v in model.state_dict().items()}
+rp=creps(model,pids_a,dev)
+p1_eval={op:ev(model,lds[f'{op}_test'],dev) for op in ALL_OPS}
+print(f"\nPhase 1 final accuracies:")
+for op,e in p1_eval.items():print(f"  {op}: {e['acc']:.3f}")
+
+os.makedirs('results',exist_ok=True)
+torch.save(model.state_dict(),'results/p1.pt')
+p1_fps = fourier_power_spectrum(model.te.weight.detach().cpu(), p)
+
+# ========== PHASE 2: Fork into 5 branches ==========
+branch_results = {}
+phase2_epochs = 150
+
+for task_name in ALL_OPS:
+    branch_name = f'p2_{task_name}'
+    level = ALL_OPS.index(task_name)
+    print(f"\n{'='*60}")
+    print(f"PHASE 2 — Branch A→{task_name.upper()} (Level {level}, {phase2_epochs} epochs)")
+    print(f"{'='*60}")
+
+    m = ST(cfg).to(dev)
+    m.load_state_dict(torch.load('results/p1.pt', weights_only=True))
+    o = optim.AdamW(m.parameters(), lr=1e-3, weight_decay=1.0)
+    h = train_phase(m, o, lds[f'{task_name}_train'], phase2_epochs, dev,
+                    branch_name, s0, s1, ri, rp, pids_a, lds, cfg, ce=20)
+
+    # Final evaluation on ALL tasks
+    final_eval = {op: ev(m, lds[f'{op}_test'], dev) for op in ALL_OPS}
+    fps_final = fourier_power_spectrum(m.te.weight.detach().cpu(), p)
+
+    branch_results[task_name] = {
+        'history': h,
+        'final_eval': final_eval,
+        'fourier_spectrum_final': fps_final.tolist(),
+    }
+
+    print(f"\n  Final accuracies after A→{task_name}:")
+    for op, e in final_eval.items():
+        marker = " ←TRAINED" if op == task_name else (" ←BASE" if op == 'add' else "")
+        print(f"    {op}: {e['acc']:.3f}{marker}")
+
+    # Forgetting = Phase1 add accuracy - current add accuracy
+    forgetting = p1_eval['add']['acc'] - final_eval['add']['acc']
+    print(f"  Addition FORGETTING: {forgetting*100:.1f}%")
+
+# ========== PHASE 3: Cross-comparison ==========
+print(f"\n{'='*60}\nPHASE 3: Cross-model Representation Comparison\n{'='*60}")
+
+# Reload all branch final models for comparison
+branch_reps = {}
+for task_name in ALL_OPS:
+    # Re-run the final model to get representations
+    # (models were not saved to disk, so re-extract from branch_results last step)
+    # Actually let's just use the last checkpoint reps from training
+    pass
+
+# Build summary table
+summary = {
+    'config': {'p':p,'n_layers':2,'d_model':128,'n_heads':4,'d_mlp':512,
+               'phase1_epochs':150,'phase2_epochs':phase2_epochs,
+               'lr':1e-3,'weight_decay':1.0,'batch_size':512,
+               'train_frac':0.5,'seed':seed,'n_parameters':model.nparams()},
+    'dissimilarity_ladder': {
+        'add': {'level':0, 'description':'Identical (Fourier circuit)'},
+        'subtract': {'level':1, 'description':'Same Fourier circuit (sign flip)'},
+        'multiply': {'level':2, 'description':'Discrete-log Fourier circuit'},
+        'max': {'level':3, 'description':'Linear/ordinal circuit'},
+        'xor': {'level':4, 'description':'Bit-level circuit'},
+    },
+    'phase1_history': h1,
+    'phase1_final_eval': {op: p1_eval[op] for op in ALL_OPS},
+    'fourier_spectrum_init': init_fps.tolist(),
+    'fourier_spectrum_phase1': p1_fps.tolist(),
+}
+
+# Per-branch results
+for task_name, br in branch_results.items():
+    level = ALL_OPS.index(task_name)
+    summary[f'branch_{task_name}'] = {
+        'level': level,
+        'history': br['history'],
+        'final_eval': br['final_eval'],
+        'fourier_spectrum': br['fourier_spectrum_final'],
+        'addition_forgetting': p1_eval['add']['acc'] - br['final_eval']['add']['acc'],
+    }
+
+# Forgetting summary table
+print("\n" + "="*60)
+print("FORGETTING SUMMARY")
+print("="*60)
+print(f"{'Task':<15} {'Level':<8} {'Add Acc':<10} {'Task Acc':<10} {'Forgetting':<12} {'Grad Align':<12}")
+print("-"*67)
+for task_name in ALL_OPS:
+    br = branch_results[task_name]
+    level = ALL_OPS.index(task_name)
+    add_acc = br['final_eval']['add']['acc']
+    task_acc = br['final_eval'][task_name]['acc']
+    forgetting = p1_eval['add']['acc'] - add_acc
+    # Get last gradient alignment from history
+    last_ga = 0
+    for h in reversed(br['history']):
+        if f'grad_align_add_vs_{task_name}' in h:
+            last_ga = h[f'grad_align_add_vs_{task_name}']; break
+        elif task_name == 'add' and 'grad_align_add_vs_subtract' in h:
+            last_ga = h['grad_align_add_vs_subtract']; break
+    print(f"{task_name:<15} {level:<8} {add_acc:<10.3f} {task_acc:<10.3f} {forgetting*100:<12.1f}% {last_ga:<12.3f}")
+
+summary['forgetting_table'] = {}
+for task_name in ALL_OPS:
+    br = branch_results[task_name]
+    summary['forgetting_table'][task_name] = {
+        'level': ALL_OPS.index(task_name),
+        'addition_accuracy': br['final_eval']['add']['acc'],
+        'task_accuracy': br['final_eval'][task_name]['acc'],
+        'forgetting_pct': (p1_eval['add']['acc'] - br['final_eval']['add']['acc']) * 100,
+    }
+
+with open('results/graduated_experiment_results.json','w') as f:
+    json.dump(summary,f,indent=2,default=str)
+
+elapsed = time.time() - t0
+print(f"\nExperiment completed in {elapsed/60:.1f} minutes")
+print("Results saved to results/graduated_experiment_results.json")
+
+# ============= GENERATE PLOTS =============
+print("\nGenerating plots...")
+import matplotlib;matplotlib.use('Agg');import matplotlib.pyplot as plt
+
+colors = {'add':'#2196F3','subtract':'#4CAF50','multiply':'#FF9800','max':'#F44336','xor':'#9C27B0'}
+labels = {'add':'Addition (L0)','subtract':'Subtraction (L1)','multiply':'Multiplication (L2)',
+          'max':'Max (L3)','xor':'XOR (L4)'}
+
+# Plot 1: Forgetting bar chart
+fig,ax=plt.subplots(figsize=(10,6))
+tasks=[t for t in ALL_OPS]
+forg=[summary['forgetting_table'][t]['forgetting_pct'] for t in tasks]
+task_acc=[summary['forgetting_table'][t]['task_accuracy'] for t in tasks]
+x=np.arange(len(tasks))
+bars=ax.bar(x,forg,color=[colors[t] for t in tasks],edgecolor='black',linewidth=0.5)
+ax.set_xticks(x);ax.set_xticklabels([labels[t] for t in tasks],rotation=15,ha='right')
+ax.set_ylabel('Addition Accuracy Forgetting (%)');ax.set_title('Forgetting vs Task Dissimilarity Level')
+ax.axhline(0,color='gray',ls=':',alpha=0.5)
+for bar,ta in zip(bars,task_acc):
+    ax.text(bar.get_x()+bar.get_width()/2,bar.get_height()+0.5,f'Task:{ta:.1%}',
+            ha='center',va='bottom',fontsize=8)
+plt.tight_layout();plt.savefig('results/forgetting_ladder.png',dpi=150);plt.close()
+
+# Plot 2: Addition accuracy over training for each branch
+fig,ax=plt.subplots(figsize=(12,6))
+s1s=[h['step'] for h in h1]
+ax.plot(s1s,[h.get('eval/add_test_acc',0) for h in h1],'k-',lw=2,label='Phase 1 (training)')
+for task_name in ALL_OPS:
+    br=branch_results[task_name]
+    steps=[h['step']+s1s[-1] for h in br['history']]
+    acc=[h.get('eval/add_test_acc',0) for h in br['history']]
+    ax.plot(steps,acc,'-',color=colors[task_name],lw=2,label=labels[task_name])
+ax.axvline(s1s[-1],color='gray',ls='--',alpha=0.5,label='Phase transition')
+ax.set_xlabel('Training Step');ax.set_ylabel('Addition Test Accuracy')
+ax.set_title('Addition Accuracy: Learning vs Forgetting');ax.legend();ax.set_ylim(-0.05,1.05)
+plt.tight_layout();plt.savefig('results/addition_accuracy_all_branches.png',dpi=150);plt.close()
+
+# Plot 3: CKA vs Phase 1 for each branch (layer 2 = final layer)
+fig,axes=plt.subplots(1,3,figsize=(18,5))
+for li,ax in enumerate(axes):
+    m_name=f'layer_{li}/cka_vs_phase1'
+    for task_name in ALL_OPS:
+        br=branch_results[task_name]
+        steps=[h['step'] for h in br['history'] if m_name in h]
+        vals=[h[m_name] for h in br['history'] if m_name in h]
+        ax.plot(steps,vals,'-',color=colors[task_name],lw=1.5,label=labels[task_name])
+    ax.set_title(f'{"Embedding" if li==0 else f"Layer {li}"}: CKA vs Phase 1 End')
+    ax.set_xlabel('Step');ax.set_ylabel('CKA');ax.legend(fontsize=7);ax.set_ylim(0,1.05)
+plt.tight_layout();plt.savefig('results/cka_all_branches.png',dpi=150);plt.close()
+
+# Plot 4: Gradient alignment with addition over training
+fig,ax=plt.subplots(figsize=(12,6))
+for task_name in ['subtract','multiply','max','xor']:
+    m_name=f'grad_align_add_vs_{task_name}'
+    # Phase 1 data
+    steps_p1=[h['step'] for h in h1 if m_name in h]
+    vals_p1=[h[m_name] for h in h1 if m_name in h]
+    ax.plot(steps_p1,vals_p1,':',color=colors[task_name],lw=1,alpha=0.5)
+    # Branch data (training on this task)
+    br=branch_results[task_name]
+    steps=[h['step']+s1s[-1] for h in br['history'] if m_name in h]
+    vals=[h[m_name] for h in br['history'] if m_name in h]
+    ax.plot(steps,vals,'-',color=colors[task_name],lw=2,label=labels[task_name])
+ax.axvline(s1s[-1],color='gray',ls='--',alpha=0.5)
+ax.axhline(0,color='gray',ls=':',alpha=0.3)
+ax.set_xlabel('Step');ax.set_ylabel('Gradient Cosine Similarity')
+ax.set_title('Gradient Alignment with Addition Task');ax.legend()
+plt.tight_layout();plt.savefig('results/gradient_alignment_all.png',dpi=150);plt.close()
+
+# Plot 5: Fourier spectrum comparison
+fig,axes=plt.subplots(2,3,figsize=(18,10))
+freqs=np.arange(len(init_fps))
+axes[0,0].bar(freqs,init_fps,color='gray',alpha=0.5);axes[0,0].set_title('Init');axes[0,0].set_ylabel('Power')
+axes[0,1].bar(freqs,p1_fps,color='black',alpha=0.7);axes[0,1].set_title('After Phase 1 (Addition)')
+for i,task_name in enumerate(['subtract','multiply','max','xor']):
+    ax=axes[(i+2)//3,(i+2)%3]
+    fps=np.array(branch_results[task_name]['fourier_spectrum_final'])
+    ax.bar(freqs[:len(fps)],fps,color=colors[task_name],alpha=0.7)
+    ax.set_title(f'After A→{task_name.title()}')
+    if (i+2)//3==1:ax.set_xlabel('Frequency')
+    if (i+2)%3==0:ax.set_ylabel('Power')
+plt.suptitle('Embedding Fourier Power Spectrum Evolution',fontsize=14)
+plt.tight_layout();plt.savefig('results/fourier_spectra.png',dpi=150);plt.close()
+
+# Plot 6: Subspace angles
+fig,ax=plt.subplots(figsize=(12,6))
+for task_name in ALL_OPS:
+    m_name='layer_2/subspace_angle_vs_phase1'
+    br=branch_results[task_name]
+    steps=[h['step'] for h in br['history'] if m_name in h]
+    vals=[h[m_name] for h in br['history'] if m_name in h]
+    ax.plot(steps,vals,'-',color=colors[task_name],lw=2,label=labels[task_name])
+ax.set_xlabel('Step');ax.set_ylabel('Subspace Angle (degrees)')
+ax.set_title('Final Layer Subspace Angle Drift from Phase 1');ax.legend()
+plt.tight_layout();plt.savefig('results/subspace_angles_all.png',dpi=150);plt.close()
+
+print("All plots generated!")
+
+# Upload to HF
+from huggingface_hub import HfApi
+api=HfApi(token=os.environ.get('HF_TOKEN'))
+repo='tekkmaven/representation-learning-dynamics'
+files_to_upload = [
+    'graduated_experiment_results.json',
+    'forgetting_ladder.png','addition_accuracy_all_branches.png',
+    'cka_all_branches.png','gradient_alignment_all.png',
+    'fourier_spectra.png','subspace_angles_all.png',
+]
+for f in files_to_upload:
+    try:
+        api.upload_file(path_or_fileobj=f'results/{f}',path_in_repo=f'results/{f}',repo_id=repo,repo_type='model')
+        print(f"Uploaded {f}")
+    except Exception as e:
+        print(f"Failed to upload {f}: {e}")
+
+print(f"\n=== EXPERIMENT COMPLETE ===")
+print(f"Total time: {(time.time()-t0)/60:.1f} minutes")
+print(f"Repository: https://huggingface.co/{repo}")