Create experiment_3_dual_teacher_autograd.py

experiment_2/experiment_3_dual_teacher_autograd.py

@@ -0,0 +1,689 @@
+# ============================================================================
+# DUAL-TEACHER PROCRUSTES CONSENSUS DISTILLATION
+#
+# Teacher A: Raw Adam (0.626) — learned without geometric guidance
+# Teacher B: Geometric (+spr+ort, 0.645) — learned with manifold control
+#
+# Both teachers encode the full training set.
+# Procrustes alignment finds the shared geometric center.
+# Student distills the consensus + trains with entropy config.
+#
+# Same pipeline as CaptionBERT 5-expert consensus, simplified to 2.
+# ============================================================================
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+print("=" * 65)
+print("DUAL-TEACHER PROCRUSTES CONSENSUS DISTILLATION")
+print("=" * 65)
+print(f"  Device: {DEVICE}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# GEOMETRIC PRIMITIVES
+# ══════════════════════════════════════════════════════════════════
+
+def tangential_projection(grad, embedding):
+    emb_n = F.normalize(embedding.detach().float(), dim=-1)
+    grad_f = grad.float()
+    radial = (grad_f * emb_n).sum(dim=-1, keepdim=True) * emb_n
+    return (grad_f - radial).to(grad.dtype), radial.to(grad.dtype)
+
+def cayley_menger_vol2(pts):
+    pts = pts.float()
+    diff = pts.unsqueeze(-2) - pts.unsqueeze(-3)
+    d2 = (diff * diff).sum(-1)
+    B, V, _ = d2.shape
+    cm = torch.zeros(B, V+1, V+1, device=d2.device, dtype=torch.float32)
+    cm[:, 0, 1:] = 1; cm[:, 1:, 0] = 1; cm[:, 1:, 1:] = d2
+    s = (-1.0)**V; f = math.factorial(V-1)
+    return s / ((2.0**(V-1)) * f*f) * torch.linalg.det(cm)
+
+def cv_loss(emb, target=0.2, n_samples=16):
+    B = emb.shape[0]
+    if B < 5: return torch.tensor(0.0, device=emb.device)
+    vols = []
+    for _ in range(n_samples):
+        idx = torch.randperm(B, device=emb.device)[:5]
+        v2 = cayley_menger_vol2(emb[idx].unsqueeze(0))
+        vols.append(torch.sqrt(F.relu(v2[0]) + 1e-12))
+    stacked = torch.stack(vols)
+    cv = stacked.std() / (stacked.mean() + 1e-8)
+    return (cv - target).abs()
+
+@torch.no_grad()
+def cv_metric(emb, n_samples=200):
+    B = emb.shape[0]
+    if B < 5: return 0.0
+    emb_f = emb.detach().float()
+    vols = []
+    for _ in range(n_samples):
+        idx = torch.randperm(B, device=emb.device)[:5]
+        v2 = cayley_menger_vol2(emb_f[idx].unsqueeze(0))
+        v = torch.sqrt(F.relu(v2[0]) + 1e-12).item()
+        if v > 0: vols.append(v)
+    if len(vols) < 10: return 0.0
+    a = torch.tensor(vols)
+    return float(a.std() / (a.mean() + 1e-8))
+
+def anchor_spread_loss(anchors):
+    a_n = F.normalize(anchors, dim=-1)
+    sim = a_n @ a_n.T - torch.diag(torch.ones(anchors.shape[0], device=anchors.device))
+    return sim.pow(2).mean()
+
+def anchor_entropy_loss(emb, anchors, sharpness=10.0):
+    a_n = F.normalize(anchors, dim=-1)
+    probs = F.softmax(emb @ a_n.T * sharpness, dim=-1)
+    return -(probs * (probs + 1e-12).log()).sum(-1).mean()
+
+def anchor_ortho_loss(anchors):
+    a_n = F.normalize(anchors, dim=-1)
+    gram = a_n @ a_n.T
+    N = anchors.shape[0]
+    mask = ~torch.eye(N, dtype=bool, device=anchors.device)
+    return gram[mask].pow(2).mean()
+
+def infonce(a, b, temperature=0.07):
+    a = F.normalize(a, dim=-1)
+    b = F.normalize(b, dim=-1)
+    logits = (a @ b.T) / temperature
+    labels = torch.arange(logits.shape[0], device=logits.device)
+    loss = (F.cross_entropy(logits, labels) + F.cross_entropy(logits.T, labels)) / 2
+    with torch.no_grad():
+        acc = (logits.argmax(-1) == labels).float().mean().item()
+    return loss, acc
+
+
+# ── Backward filtering ──
+
+class EmbeddingAutograd(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, embedding, anchors, tang, sep):
+        ctx.save_for_backward(embedding, anchors)
+        ctx.tang = tang; ctx.sep = sep
+        return x
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        embedding, anchors = ctx.saved_tensors
+        emb_n = F.normalize(embedding.detach().float(), dim=-1)
+        anchors_n = F.normalize(anchors.detach().float(), dim=-1)
+        grad_f = grad_output.float()
+        tang_grad, norm_grad = tangential_projection(grad_f, emb_n)
+        corrected = tang_grad + (1.0 - ctx.tang) * norm_grad
+        if ctx.sep > 0:
+            cos_to = emb_n @ anchors_n.T
+            nearest = anchors_n[cos_to.argmax(dim=-1)]
+            toward = (corrected * nearest).sum(dim=-1, keepdim=True)
+            collapse = toward * nearest
+            corrected = corrected - ctx.sep * (toward > 0).float() * collapse
+        return corrected.to(grad_output.dtype), None, None, None, None
+
+
+# ════════════��═════════════════════════════════════════════════════
+# PROCRUSTES ALIGNMENT (production, from cotrain_bank.py)
+# ══════════════════════════════════════════════════════════════════
+
+def symmetric_inv_sqrt(cov, eps=1e-6):
+    evals, evecs = torch.linalg.eigh(cov)
+    evals = torch.clamp(evals, min=eps)
+    return evecs @ torch.diag(evals.rsqrt()) @ evecs.T
+
+def procrustes_align(source, target, n_align=10000):
+    N = min(n_align, source.shape[0], target.shape[0])
+    S = source[:N].float(); T = target[:N].float()
+    s_mean = S.mean(0, keepdim=True); t_mean = T.mean(0, keepdim=True)
+    Sc = S - s_mean; Tc = T - t_mean; Ns = Sc.shape[0]
+    s_cov = (Sc.T @ Sc) / max(Ns - 1, 1)
+    t_cov = (Tc.T @ Tc) / max(Ns - 1, 1)
+    s_whiten = symmetric_inv_sqrt(s_cov)
+    t_whiten = symmetric_inv_sqrt(t_cov)
+    Sc_w = F.normalize(Sc @ s_whiten, dim=-1)
+    Tc_w = F.normalize(Tc @ t_whiten, dim=-1)
+    U, _, Vt = torch.linalg.svd(Tc_w.T @ Sc_w, full_matrices=False)
+    R = U @ Vt
+    cos_after = F.cosine_similarity(Sc_w @ R.T, Tc_w, dim=-1).mean().item()
+    return {"rotation": R, "source_mean": s_mean.squeeze(0),
+            "source_whitener": s_whiten, "cos_after": cos_after}
+
+def apply_align(emb, info):
+    x = emb.float() - info["source_mean"]
+    x = x @ info["source_whitener"]
+    x = x @ info["rotation"].T
+    return x
+
+
+# ══════════════════════════════════════════════════════════════════
+# MODEL
+# ══════════════════════════════════════════════════════════════════
+
+class Constellation(nn.Module):
+    def __init__(self, n_anchors=30, d_embed=768, init_anchors=None):
+        super().__init__()
+        self.n_anchors = n_anchors
+        if init_anchors is not None:
+            self.anchors = nn.Parameter(init_anchors.clone())
+        else:
+            self.anchors = nn.Parameter(F.normalize(torch.randn(n_anchors, d_embed), dim=-1))
+        self.register_buffer("rigidity", torch.zeros(n_anchors))
+        self.register_buffer("visit_count", torch.zeros(n_anchors))
+
+    def triangulate(self, emb):
+        anchors_n = F.normalize(self.anchors, dim=-1)
+        cos_sim = emb @ anchors_n.T
+        return 1.0 - cos_sim, cos_sim.argmax(dim=-1)
+
+    @torch.no_grad()
+    def update_rigidity(self, tri_dist):
+        nearest = tri_dist.argmin(dim=-1)
+        for i in range(self.n_anchors):
+            mask = nearest == i
+            if mask.sum() < 5: continue
+            self.visit_count[i] += mask.sum().float()
+            spread = tri_dist[mask].std(dim=0).mean()
+            alpha = min(0.1, 10.0 / (self.visit_count[i] + 1))
+            self.rigidity[i] = (1 - alpha) * self.rigidity[i] + alpha / (spread + 0.01)
+
+    @torch.no_grad()
+    def geometry_snapshot(self):
+        a_n = F.normalize(self.anchors.detach(), dim=-1)
+        cos = a_n @ a_n.T
+        mask = ~torch.eye(self.n_anchors, dtype=bool, device=a_n.device)
+        return {"anchors": a_n.clone(), "rigidity": self.rigidity.clone(),
+                "mean_cos": cos[mask].mean().item(), "std_cos": cos[mask].std().item()}
+
+
+class Patchwork(nn.Module):
+    def __init__(self, n_anchors=30, n_compartments=6, d_comp=64):
+        super().__init__()
+        self.n_compartments = n_compartments
+        assignments = torch.arange(n_anchors) % n_compartments
+        self.register_buffer("assignments", assignments)
+        self.compartments = nn.ModuleList()
+        for k in range(n_compartments):
+            n_k = (assignments == k).sum().item()
+            self.compartments.append(nn.Sequential(
+                nn.Linear(n_k, d_comp * 2), nn.GELU(),
+                nn.Linear(d_comp * 2, d_comp), nn.LayerNorm(d_comp)))
+
+    def forward(self, tri_dist):
+        parts = []
+        for k in range(self.n_compartments):
+            parts.append(self.compartments[k](tri_dist[:, self.assignments == k]))
+        return torch.cat(parts, dim=-1)
+
+
+class PatchworkClassifier(nn.Module):
+    def __init__(self, n_classes=30, n_anchors=30, d_embed=768,
+                 n_compartments=6, d_comp=64, d_hidden=256, init_anchors=None):
+        super().__init__()
+        self.backbone = nn.Sequential(
+            nn.Conv2d(1, 32, 3, padding=1), nn.GELU(), nn.MaxPool2d(2),
+            nn.Conv2d(32, 64, 3, padding=1), nn.GELU(), nn.MaxPool2d(2),
+            nn.Conv2d(64, 128, 3, padding=1), nn.GELU(), nn.AdaptiveAvgPool2d(1))
+        self.embed_proj = nn.Sequential(nn.Linear(128, d_embed), nn.LayerNorm(d_embed))
+        self.constellation = Constellation(n_anchors, d_embed, init_anchors)
+        self.patchwork = Patchwork(n_anchors, n_compartments, d_comp)
+        pw_dim = n_compartments * d_comp
+        self.mlp = nn.Sequential(
+            nn.Linear(pw_dim, d_hidden), nn.GELU(), nn.LayerNorm(d_hidden),
+            nn.Linear(d_hidden, d_hidden), nn.GELU(), nn.LayerNorm(d_hidden),
+            nn.Linear(d_hidden, n_classes))
+
+    def forward(self, x):
+        feat = self.backbone(x).flatten(1)
+        emb = F.normalize(self.embed_proj(feat), dim=-1)
+        tri, nearest = self.constellation.triangulate(emb)
+        return self.mlp(self.patchwork(tri)), emb, tri, nearest
+
+    def encode(self, x):
+        feat = self.backbone(x).flatten(1)
+        return F.normalize(self.embed_proj(feat), dim=-1)
+
+
+# ══════════════════════════════════════════════════════════════════
+# SHAPE RENDERERS
+# ══════════════════════════════════════════════════════════════════
+
+def _d(img,x0,y0,x1,y1,t=1):
+    n=max(int(max(abs(x1-x0),abs(y1-y0))*2),1);sz=img.shape[0]
+    for s in np.linspace(0,1,n):
+        px,py=int(x0+s*(x1-x0)),int(y0+s*(y1-y0))
+        for dx in range(-t,t+1):
+            for dy in range(-t,t+1):
+                nx,ny=px+dx,py+dy
+                if 0<=nx<sz and 0<=ny<sz: img[ny,nx]=1.0
+
+def rpoly(nv,sz=32,p=0.15):
+    img=np.zeros((sz,sz),dtype=np.float32);cx,cy,r=sz/2,sz/2,sz*0.35
+    a=np.linspace(0,2*np.pi,nv,endpoint=False)+np.random.uniform(0,2*np.pi)
+    ri=r*(1+np.random.normal(0,p,nv))
+    pts=[(cx+ri[i]*np.cos(a[i]),cy+ri[i]*np.sin(a[i])) for i in range(nv)]
+    for i in range(nv): _d(img,*pts[i],*pts[(i+1)%nv])
+    return img
+
+def rstar(np_,sz=32,p=0.12):
+    img=np.zeros((sz,sz),dtype=np.float32);cx,cy=sz/2,sz/2;ro,ri_=sz*0.38,sz*0.15
+    a=np.linspace(0,2*np.pi,np_*2,endpoint=False)+np.random.uniform(0,2*np.pi)
+    pts=[(cx+(ro if i%2==0 else ri_)*(1+np.random.normal(0,p))*np.cos(a[i]),
+          cy+(ro if i%2==0 else ri_)*(1+np.random.normal(0,p))*np.sin(a[i])) for i in range(len(a))]
+    for i in range(len(pts)): _d(img,*pts[i],*pts[(i+1)%len(pts)])
+    return img
+
+def rcross(sz=32,p=0.15):
+    img=np.zeros((sz,sz),dtype=np.float32);cx,cy,arm=sz/2,sz/2,sz*0.3
+    for ab in [0,np.pi/2,np.pi,3*np.pi/2]:
+        a=ab+np.random.normal(0,p*0.3);r=arm*(1+np.random.normal(0,p))
+        _d(img,cx,cy,cx+r*np.cos(a),cy+r*np.sin(a),2)
+    return img
+
+def rspiral(sz=32,p=0.1):
+    img=np.zeros((sz,sz),dtype=np.float32);cx,cy=sz/2,sz/2
+    for t in np.linspace(0,5*np.pi,200):
+        r=sz*0.015*t*(1+np.random.normal(0,p*0.3));x,y=int(cx+r*np.cos(t)),int(cy+r*np.sin(t))
+        if 0<=x<sz and 0<=y<sz: img[y,x]=1.0
+    return img
+
+def rwave(sz=32,p=0.1):
+    img=np.zeros((sz,sz),dtype=np.float32);f=2+np.random.normal(0,0.3);amp=sz*0.15*(1+np.random.normal(0,p))
+    for x in range(sz):
+        y=int(sz/2+amp*np.sin(2*np.pi*f*x/sz))
+        if 0<=y<sz: img[y,x]=1.0
+    return img
+
+def rheart(sz=32,p=0.1):
+    img=np.zeros((sz,sz),dtype=np.float32);cx,cy=sz/2,sz*0.45;s=sz*0.017*(1+np.random.normal(0,p))
+    for t in np.linspace(0,2*np.pi,300):
+        x=16*np.sin(t)**3;y=-(13*np.cos(t)-5*np.cos(2*t)-2*np.cos(3*t)-np.cos(4*t))
+        ix,iy=int(cx+x*s),int(cy+y*s)
+        if 0<=ix<sz and 0<=iy<sz: img[iy,ix]=1.0
+    return img
+
+def rcrescent(sz=32,p=0.1):
+    img=np.zeros((sz,sz),dtype=np.float32);cx,cy,r=sz/2,sz/2,sz*0.35;r2=r*0.7;off=r*0.3
+    for a in np.linspace(0,2*np.pi,300):
+        x1,y1=cx+r*np.cos(a),cy+r*np.sin(a)
+        if math.sqrt((x1-cx-off)**2+(y1-cy)**2)>=r2*0.9:
+            ix,iy=int(x1),int(y1)
+            if 0<=ix<sz and 0<=iy<sz: img[iy,ix]=1.0
+    return img
+
+def rellipse(sz=32,p=0.1):
+    img=np.zeros((sz,sz),dtype=np.float32);cx,cy=sz/2,sz/2
+    a,b=sz*0.38*(1+np.random.normal(0,p)),sz*0.22*(1+np.random.normal(0,p));rot=np.random.uniform(0,np.pi)
+    for t in np.linspace(0,2*np.pi,200):
+        x,y=a*np.cos(t),b*np.sin(t);ix,iy=int(cx+x*np.cos(rot)-y*np.sin(rot)),int(cy+x*np.sin(rot)+y*np.cos(rot))
+        if 0<=ix<sz and 0<=iy<sz: img[iy,ix]=1.0
+    return img
+
+def rring(sz=32,p=0.1):
+    img=np.zeros((sz,sz),dtype=np.float32);cx,cy=sz/2,sz/2
+    r1,r2=sz*0.35*(1+np.random.normal(0,p)),sz*0.22*(1+np.random.normal(0,p))
+    for a in np.linspace(0,2*np.pi,300):
+        for r in [r1,r2]:
+            x,y=int(cx+r*np.cos(a)),int(cy+r*np.sin(a))
+            if 0<=x<sz and 0<=y<sz: img[y,x]=1.0
+    return img
+
+def rarrow(sz=32,p=0.12):
+    img=np.zeros((sz,sz),dtype=np.float32);cx,cy=sz/2,sz/2
+    l=sz*0.35*(1+np.random.normal(0,p));h=l*0.35;a=np.random.uniform(0,2*np.pi)
+    x1,y1=cx-l*np.cos(a),cy-l*np.sin(a);x2,y2=cx+l*np.cos(a),cy+l*np.sin(a)
+    _d(img,x1,y1,x2,y2)
+    for da in [0.7,-0.7]: _d(img,x2,y2,x2-h*np.cos(a+da),y2-h*np.sin(a+da))
+    return img
+
+def rchevron(sz=32,p=0.12):
+    img=np.zeros((sz,sz),dtype=np.float32);cx,cy=sz/2,sz/2
+    w,h=sz*0.3*(1+np.random.normal(0,p)),sz*0.25*(1+np.random.normal(0,p))
+    _d(img,cx-w,cy+h,cx,cy-h);_d(img,cx,cy-h,cx+w,cy+h)
+    return img
+
+def rsemicirc(sz=32,p=0.1):
+    img=np.zeros((sz,sz),dtype=np.float32);cx,cy,r=sz/2,sz*0.6,sz*0.35
+    for a in np.linspace(np.pi,2*np.pi,150):
+        x,y=int(cx+r*np.cos(a)),int(cy+r*np.sin(a))
+        if 0<=x<sz and 0<=y<sz: img[y,x]=1.0
+    _d(img,cx-r,cy,cx+r,cy)
+    return img
+
+def gen_one(c,sz=32):
+    if c==0: return rpoly(3,sz,0.20)
+    if c==1: return rpoly(4,sz,0.12)
+    if c==2: return rpoly(5,sz,0.15)
+    if c==3: return rpoly(6,sz,0.10)
+    if c==4: return rpoly(7,sz,0.10)
+    if c==5: return rpoly(8,sz,0.08)
+    if c==6: return rpoly(9,sz,0.08)
+    if c==7: return rpoly(10,sz,0.07)
+    if c==8: return rpoly(12,sz,0.06)
+    if c==9: return rpoly(32,sz,0.03)
+    if c==10: return rellipse(sz)
+    if c==11: return rspiral(sz)
+    if c==12: return rwave(sz)
+    if c==13: return rcrescent(sz)
+    if c==14: return rstar(3,sz)
+    if c==15: return rstar(4,sz)
+    if c==16: return rstar(5,sz)
+    if c==17: return rstar(6,sz)
+    if c==18: return rstar(7,sz)
+    if c==19: return rstar(8,sz)
+    if c==20: return rcross(sz)
+    if c==21: return rpoly(4,sz,0.10)
+    if c==22: return rarrow(sz)
+    if c==23: return rheart(sz)
+    if c==24: return rring(sz)
+    if c==25: return rsemicirc(sz)
+    if c==26: return rpoly(4,sz,0.15)
+    if c==27: return rpoly(4,sz,0.18)
+    if c==28: return rpoly(4,sz,0.10)
+    if c==29: return rchevron(sz)
+    return rpoly(3,sz)
+
+def gen_data(n_per=500, sz=32):
+    imgs, labels = [], []
+    for _ in range(n_per):
+        for c in range(30):
+            imgs.append(gen_one(c, sz)); labels.append(c)
+    imgs = torch.tensor(np.array(imgs)).unsqueeze(1)
+    labels = torch.tensor(labels, dtype=torch.long)
+    perm = torch.randperm(len(labels))
+    return imgs[perm], labels[perm]
+
+
+# ══════════════════════════════════════════════════════════════════
+# GENERATE SHARED DATA
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n  Generating data...")
+torch.manual_seed(42); np.random.seed(42)
+train_imgs, train_labels = gen_data(n_per=500)
+val_imgs, val_labels = gen_data(n_per=100)
+train_imgs, train_labels = train_imgs.to(DEVICE), train_labels.to(DEVICE)
+val_imgs, val_labels = val_imgs.to(DEVICE), val_labels.to(DEVICE)
+n_train, n_val = len(train_labels), len(val_labels)
+print(f"  Train: {n_train:,}  Val: {n_val:,}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TRAIN TEACHER HELPER
+# ══════════════════════════════════════════════════════════════════
+
+def train_teacher(model, tag, use_autograd=True,
+                  tang=0.01, sep=1.0, cv_w=0.001,
+                  w_spread=0.0, w_ortho=0.0, w_entropy=0.0,
+                  epochs=30):
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+    BATCH = 256
+    for epoch in range(epochs):
+        model.train()
+        perm = torch.randperm(n_train, device=DEVICE)
+        total_correct, n = 0, 0
+        for i in range(0, n_train, BATCH):
+            idx = perm[i:i+BATCH]
+            if len(idx) < 4: continue
+            logits, emb, tri, nearest = model(train_imgs[idx])
+            labels = train_labels[idx]
+            anchors = model.constellation.anchors
+            if use_autograd and (tang > 0 or sep > 0):
+                emb_g = EmbeddingAutograd.apply(emb, emb, anchors, tang, sep)
+                tri_g, _ = model.constellation.triangulate(emb_g)
+                logits = model.mlp(model.patchwork(tri_g))
+            l_cls = F.cross_entropy(logits, labels)
+            l_geo = torch.tensor(0.0, device=DEVICE)
+            if cv_w > 0: l_geo = l_geo + cv_w * cv_loss(emb)
+            if w_spread > 0: l_geo = l_geo + w_spread * anchor_spread_loss(anchors)
+            if w_ortho > 0: l_geo = l_geo + w_ortho * anchor_ortho_loss(anchors)
+            if w_entropy > 0: l_geo = l_geo + w_entropy * anchor_entropy_loss(emb, anchors)
+            (l_cls + l_geo).backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step(); optimizer.zero_grad(set_to_none=True)
+            model.constellation.update_rigidity(tri.detach())
+            total_correct += (logits.argmax(-1) == labels).sum().item()
+            n += 1
+        model.eval()
+        with torch.no_grad():
+            vl, ve, _, _ = model(val_imgs)
+            v_acc = (vl.argmax(-1) == val_labels).float().mean().item()
+            v_cv = cv_metric(ve)
+        if (epoch+1) % 10 == 0 or epoch == 0:
+            print(f"  {tag} E{epoch+1:2d}: t={total_correct/n_train:.3f} "
+                  f"v={v_acc:.3f} cv={v_cv:.4f}")
+    return v_acc
+
+
+# ══════════════════════════════════════════════════════════════════
+# STAGE 1: TRAIN BOTH TEACHERS
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'='*65}")
+print("STAGE 1A: TEACHER A — Raw Adam")
+print(f"{'='*65}")
+torch.manual_seed(42)
+teacher_a = PatchworkClassifier(n_classes=30, n_anchors=30, d_embed=768).to(DEVICE)
+va_a = train_teacher(teacher_a, "[A]", use_autograd=False, epochs=30)
+
+print(f"\n{'='*65}")
+print("STAGE 1B: TEACHER B — Geometric (+spr+ort)")
+print(f"{'='*65}")
+torch.manual_seed(42)
+teacher_b = PatchworkClassifier(n_classes=30, n_anchors=30, d_embed=768).to(DEVICE)
+va_b = train_teacher(teacher_b, "[B]", use_autograd=True,
+                     tang=0.01, sep=1.0, cv_w=0.001,
+                     w_spread=1e-3, w_ortho=1e-3, epochs=30)
+
+
+# ══════════════════════════════════════════════════════════════════
+# STAGE 2: EXTRACT EMBEDDINGS + GPA ALIGNMENT
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'='*65}")
+print("STAGE 2: EXTRACT + PROCRUSTES ALIGN")
+print(f"{'='*65}")
+
+teacher_a.eval(); teacher_b.eval()
+with torch.no_grad():
+    emb_a = teacher_a.encode(train_imgs)  # (N, 768)
+    emb_b = teacher_b.encode(train_imgs)  # (N, 768)
+    val_emb_a = teacher_a.encode(val_imgs)
+    val_emb_b = teacher_b.encode(val_imgs)
+
+print(f"  Teacher A embeddings: {emb_a.shape}")
+print(f"  Teacher B embeddings: {emb_b.shape}")
+print(f"  Raw cos(A, B): {F.cosine_similarity(emb_a[:1000], emb_b[:1000], dim=-1).mean():.4f}")
+
+# GPA: iterative Procrustes to find geometric center
+current = {"a": emb_a.float(), "b": emb_b.float()}
+for gpa_iter in range(10):
+    mean_shape = (current["a"] + current["b"]) / 2
+    total_delta = 0.0
+    new_current = {}
+    for name in ["a", "b"]:
+        info = procrustes_align(current[name], mean_shape)
+        new_current[name] = apply_align(current[name], info)
+        total_delta += (new_current[name] - current[name]).pow(2).mean().item()
+    current = new_current
+    if (gpa_iter + 1) % 5 == 0 or gpa_iter == 0:
+        print(f"  GPA iter {gpa_iter+1}: delta={total_delta:.8f}")
+    if total_delta < 1e-8:
+        print(f"  Converged at iteration {gpa_iter+1}")
+        break
+
+# Build consensus
+mean_shape = (current["a"] + current["b"]) / 2
+consensus = F.normalize(mean_shape, dim=-1)
+
+# Align val embeddings too
+val_aligned = {}
+for name, emb_full in [("a", emb_a), ("b", emb_b)]:
+    info = procrustes_align(emb_full, mean_shape)
+    val_aligned[name] = apply_align(
+        val_emb_a if name == "a" else val_emb_b, info)
+
+val_consensus = F.normalize((val_aligned["a"] + val_aligned["b"]) / 2, dim=-1)
+
+# Check alignment quality
+for name in ["a", "b"]:
+    aligned = current[name] if name == "a" else current["b"]
+    cos = F.cosine_similarity(consensus[:1000], F.normalize(aligned[:1000], dim=-1), dim=-1).mean()
+    print(f"  cos(consensus, {name}): {cos:.4f}")
+
+consensus_cv = cv_metric(consensus[:2000])
+print(f"  Consensus CV: {consensus_cv:.4f}")
+
+# Extract consensus anchors: cluster consensus embeddings
+# Use the per-class centroids as initial anchor positions
+anchor_centroids = []
+for c in range(30):
+    mask = train_labels == c
+    if mask.sum() > 0:
+        anchor_centroids.append(consensus[mask].mean(dim=0))
+    else:
+        anchor_centroids.append(torch.randn(768, device=DEVICE))
+consensus_anchors = F.normalize(torch.stack(anchor_centroids), dim=-1)
+print(f"  Consensus anchors: {consensus_anchors.shape}")
+
+# Also keep the pure geometric center from teacher anchors
+geo_a = teacher_a.constellation.geometry_snapshot()
+geo_b = teacher_b.constellation.geometry_snapshot()
+print(f"  Teacher A anchors cos: {geo_a['mean_cos']:.4f}")
+print(f"  Teacher B anchors cos: {geo_b['mean_cos']:.4f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# STAGE 3: STUDENT DISTILLATION
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'='*65}")
+print("STAGE 3: STUDENT — Consensus distillation + classification")
+print(f"{'='*65}")
+
+torch.manual_seed(42)
+model_student = PatchworkClassifier(
+    n_classes=30, n_anchors=30, d_embed=768,
+    init_anchors=consensus_anchors,
+).to(DEVICE)
+
+optimizer = torch.optim.Adam(model_student.parameters(), lr=1e-3)
+BATCH = 256
+EPOCHS = 30
+
+# Move consensus targets to device
+train_targets = consensus.to(DEVICE)  # (N, 768) L2-normalized
+val_targets = val_consensus.to(DEVICE)
+
+for epoch in range(EPOCHS):
+    model_student.train()
+    perm = torch.randperm(n_train, device=DEVICE)
+    total_loss, total_correct, n = 0, 0, 0
+
+    for i in range(0, n_train, BATCH):
+        idx = perm[i:i+BATCH]
+        if len(idx) < 4: continue
+
+        logits, emb, tri, nearest = model_student(train_imgs[idx])
+        labels = train_labels[idx]
+        tgt = train_targets[idx]
+        anchors = model_student.constellation.anchors
+
+        # Apply geometric autograd
+        emb_g = EmbeddingAutograd.apply(emb, emb, anchors, 0.01, 1.0)
+        tri_g, _ = model_student.constellation.triangulate(emb_g)
+        logits = model_student.mlp(model_student.patchwork(tri_g))
+
+        # Three losses:
+        # 1. Classification (task)
+        l_cls = F.cross_entropy(logits, labels)
+
+        # 2. Consensus distillation (InfoNCE + MSE)
+        l_nce, nce_acc = infonce(emb, tgt)
+        l_mse = F.mse_loss(emb, tgt)
+
+        # 3. Geometric (micro CV + entropy)
+        l_cv = cv_loss(emb, target=0.2)
+        l_ent = anchor_entropy_loss(emb, anchors)
+
+        loss = l_cls + 0.5 * l_nce + 0.5 * l_mse + 0.001 * l_cv + 1e-4 * l_ent
+
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model_student.parameters(), 1.0)
+        optimizer.step(); optimizer.zero_grad(set_to_none=True)
+
+        model_student.constellation.update_rigidity(tri.detach())
+        total_correct += (logits.argmax(-1) == labels).sum().item()
+        total_loss += loss.item()
+        n += 1
+
+    train_acc = total_correct / n_train
+
+    # Validation
+    model_student.eval()
+    with torch.no_grad():
+        vl, ve, _, _ = model_student(val_imgs)
+        v_acc = (vl.argmax(-1) == val_labels).float().mean().item()
+        v_cv = cv_metric(ve)
+        v_cos = F.cosine_similarity(ve, val_targets, dim=-1).mean().item()
+
+        types = {"polygon": list(range(9)), "curve": list(range(9,14)),
+                 "star": list(range(14,20)), "structure": list(range(20,30))}
+        ta = {}
+        for tname, tids in types.items():
+            tmask = torch.zeros(n_val, dtype=bool, device=DEVICE)
+            for tid in tids: tmask |= (val_labels == tid)
+            if tmask.sum() > 0:
+                ta[tname] = (vl.argmax(-1)[tmask] == val_labels[tmask]).float().mean().item()
+
+    if (epoch+1) % 5 == 0 or epoch == 0:
+        ta_str = " ".join(f"{t}={a:.2f}" for t, a in ta.items())
+        rig = model_student.constellation.rigidity
+        print(f"  E{epoch+1:2d}: t={train_acc:.3f} v={v_acc:.3f} "
+              f"cos={v_cos:.3f} cv={v_cv:.4f} "
+              f"rig={rig.mean():.1f}/{rig.max():.1f} [{ta_str}]")
+
+
+# ══════════════════════════════════════════════════════════════════
+# FINAL COMPARISON
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'='*65}")
+print("FINAL COMPARISON")
+print(f"{'='*65}")
+
+# Get final metrics for all models
+model_student.eval(); teacher_a.eval(); teacher_b.eval()
+with torch.no_grad():
+    results = {}
+    for name, mdl in [("Teacher_A", teacher_a), ("Teacher_B", teacher_b),
+                       ("Student", model_student)]:
+        vl, ve, _, _ = mdl(val_imgs)
+        acc = (vl.argmax(-1) == val_labels).float().mean().item()
+        ta = {}
+        for tname, tids in types.items():
+            tmask = torch.zeros(n_val, dtype=bool, device=DEVICE)
+            for tid in tids: tmask |= (val_labels == tid)
+            if tmask.sum() > 0:
+                ta[tname] = (vl.argmax(-1)[tmask] == val_labels[tmask]).float().mean().item()
+        results[name] = {"acc": acc, "cv": cv_metric(ve), "types": ta}
+
+print(f"\n  {'Model':<15} {'v_acc':>6} {'cv':>7} {'poly':>5} {'curve':>5} {'star':>5} {'struct':>5}")
+print(f"  {'-'*55}")
+for name, r in results.items():
+    ta = r["types"]
+    print(f"  {name:<15} {r['acc']:>6.3f} {r['cv']:>7.4f} "
+          f"{ta.get('polygon',0):>5.2f} {ta.get('curve',0):>5.2f} "
+          f"{ta.get('star',0):>5.2f} {ta.get('structure',0):>5.2f}")
+
+# Anchor drift from consensus
+s_anchors = F.normalize(model_student.constellation.anchors.detach(), dim=-1)
+drift = 1.0 - F.cosine_similarity(consensus_anchors, s_anchors, dim=-1)
+print(f"\n  Student anchor drift from consensus: mean={drift.mean():.4f} max={drift.max():.4f}")
+
+print(f"\n{'='*65}")
+print("DONE")
+print(f"{'='*65}")