Create probe.py

code/probe.py

@@ -0,0 +1,524 @@
+"""
+Comprehensive testing cell for BaselineViT (RoseFace-aware)
+Run AFTER loading your model & checkpoint in Colab.
+Assumes: model, get_cifar100_dataloaders are defined.
+"""
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+import matplotlib.pyplot as plt
+from sklearn.manifold import TSNE
+from sklearn.decomposition import PCA
+from sklearn.cluster import KMeans
+from sklearn.metrics import silhouette_score
+from pathlib import Path
+import json
+from tqdm import tqdm
+
+# =========================
+# RoseFace-aware utilities
+# =========================
+
+@torch.no_grad()
+def principal_angle_overlap(class_pentachora):
+    """
+    Measure subspace overlap between classes (lower = better decoupling).
+    Returns (mean_fro, std_fro) across all class pairs.
+    """
+    device = class_pentachora[0].vertices.device
+    dtype = class_pentachora[0].vertices.dtype
+    C = len(class_pentachora)
+    U = []
+    for p in class_pentachora:
+        V = p.vertices.to(device=device, dtype=dtype)  # [5,D]
+        c = V.mean(dim=0, keepdim=True)
+        A = V - c                                      # [5,D]
+        # QR on D x 5 (A^T) → orthonormal basis in R^D
+        Q, _ = torch.linalg.qr(A.t(), mode='reduced')  # [D, r]
+        U.append(Q)
+    overlaps = []
+    for a in range(C):
+        for b in range(a+1, C):
+            M = U[a].t() @ U[b]                        # [r_a, r_b]
+            overlaps.append(torch.linalg.norm(M, 'fro').item())
+    if not overlaps:
+        return 0.0, 0.0
+    return float(np.mean(overlaps)), float(np.std(overlaps))
+
+@torch.no_grad()
+def face_usage_heatmap(model, features_proj, targets, norm_type='l1'):
+    """
+    Compute per-class face (triad) usage heatmap [C,10].
+    features_proj: [N,D] L1-normalized (from model forward outputs)
+    """
+    device, dtype = features_proj.device, features_proj.dtype
+    C = model.num_classes
+    triplets = torch.tensor([
+        [0,1,2],[0,1,3],[0,1,4],
+        [0,2,3],[0,2,4],[0,3,4],
+        [1,2,3],[1,2,4],[1,3,4],
+        [2,3,4]
+    ], device=device, dtype=torch.long)
+    counts = torch.zeros(C, 10, device=device, dtype=torch.long)
+
+    for cls in torch.unique(targets):
+        idx = (targets == cls)
+        if idx.sum() == 0:
+            continue
+        f = features_proj[idx]  # [b,D]
+        p = model.class_pentachora[int(cls)]
+        Vn = p.vertices_norm if norm_type == 'l1' else F.normalize(p.vertices, dim=-1)  # [5,D]
+
+        # Build 10 faces
+        faces = []
+        for t in triplets:
+            b = (Vn[t[0]] + Vn[t[1]] + Vn[t[2]]) / 3.0
+            if norm_type == 'l1':
+                b = b / (b.abs().sum() + 1e-8)
+            else:
+                b = F.normalize(b.unsqueeze(0), dim=-1).squeeze(0)
+            faces.append(b)
+        F10 = torch.stack(faces, dim=0)                      # [10,D]
+        sims = f @ F10.t()                                   # [b,10]
+        winner = sims.argmax(dim=1)                          # [b]
+        binc = torch.bincount(winner, minlength=10)          # [10]
+        counts[int(cls)] += binc
+
+    counts = counts.float()
+    counts = counts / (counts.sum(dim=1, keepdim=True) + 1e-9)
+    return counts  # [C,10]
+
+@torch.no_grad()
+def margin_stats(cos_pre, targets):
+    """
+    Compute margin Δ = pos - best_neg from PRE-margin cosines.
+    """
+    pos = cos_pre.gather(1, targets.view(-1,1)).squeeze(1)   # [B]
+    masked = cos_pre.masked_fill(F.one_hot(targets, cos_pre.size(1)).bool(), -1e9)
+    neg = masked.max(dim=1).values                            # [B]
+    delta = pos - neg
+    return float(delta.mean()), float(delta.std())
+
+# ============================================
+# FEATURE EXTRACTION AND ANALYSIS (upgraded)
+# ============================================
+
+class FeatureAnalyzer:
+    """
+    Analyze feature capacity and geometric patterns.
+    Now aware of RoseFace:
+      - can run with or without margin at inference (margin_mode)
+      - stores pre-margin cosines and post-margin cosines
+    """
+    def __init__(self, model, dataloader, device=None, margin_mode='none'):
+        """
+        margin_mode:
+          'none'  -> don't pass targets to model (no margin at eval)
+          'apply' -> pass targets (apply margin at eval)
+          'both'  -> run both (twice); store *_nomargin and *_margin
+        """
+        self.model = model
+        self.dataloader = dataloader
+        self.device = device or next(model.parameters()).device
+        self.model.eval()
+        assert margin_mode in ('none','apply','both')
+        self.margin_mode = margin_mode
+
+    def _forward_once(self, images, labels, apply_margin):
+        # forward; return dict of tensors on CPU
+        if apply_margin:
+            outputs = self.model(images, return_features=True, targets=labels)
+        else:
+            outputs = self.model(images, return_features=True)
+
+        out = {k: v.detach().cpu() for k, v in outputs.items() if isinstance(v, torch.Tensor)}
+        # Derive post-margin cosines (if RoseFace): cos_post = logits / s
+        if getattr(self.model, 'head_type', 'legacy') == 'roseface':
+            s = float(getattr(self.model, 'scale_s', 1.0))
+            if s > 0 and 'logits' in out:
+                out['cos_post'] = (out['logits'] / s)
+        return out
+
+    def extract_all_features(self, max_batches=None):
+        """
+        Extract features, pre-margin cosines, post-margin cosines (if available).
+        Returns dict with keys:
+          - cls_features
+          - features_proj
+          - similarities        (pre-margin cos)
+          - cos_post            (post-margin cos; RoseFace only)
+          - logits
+          - labels
+        If margin_mode == 'both', suffix *_nomargin / *_margin are included.
+        """
+        agg = {}
+
+        def append_batch(prefix, out_tensors, labels):
+            # initialize lists
+            for k, v in out_tensors.items():
+                agg.setdefault(f'{prefix}{k}', []).append(v)
+            agg.setdefault(f'{prefix}labels', []).append(labels.cpu())
+
+        with torch.no_grad():
+            for i, (images, labels) in enumerate(tqdm(self.dataloader, desc="Extracting features")):
+                if max_batches is not None and i >= max_batches:
+                    break
+                images = images.to(self.device, non_blocking=True)
+                labels = labels.to(self.device, non_blocking=True)
+
+                if self.margin_mode in ('none', 'both'):
+                    out0 = self._forward_once(images, labels, apply_margin=False)
+                    append_batch('', out0, labels)
+
+                if self.margin_mode in ('apply', 'both'):
+                    out1 = self._forward_once(images, labels, apply_margin=True)
+                    append_batch('m_', out1, labels)
+
+        # concat
+        # concat everything we collected
+        for k, lst in agg.items():
+            agg[k] = torch.cat(lst, dim=0)
+
+        # helper: pick normal key, else 'm_' fallback
+        def pick(key: str):
+            return agg.get(key, agg.get(f"m_{key}", torch.empty(0)))
+
+        # unify view for downstream code
+        if self.margin_mode == 'both':
+            features = {
+                'cls_features':   pick('features'),
+                'features_proj':  pick('features_proj'),
+                'similarities':   agg.get('similarities', torch.empty(0)),
+                'cos_post':       agg.get('cos_post', torch.empty(0)),
+                'labels':         agg.get('labels', torch.empty(0)),
+                'similarities_margin': agg.get('m_similarities', torch.empty(0)),
+                'cos_post_margin':     agg.get('m_cos_post', torch.empty(0)),
+                'logits':              agg.get('logits', torch.empty(0)),
+                'logits_margin':       agg.get('m_logits', torch.empty(0)),
+            }
+        else:
+            # works for BOTH margin_mode='none' and margin_mode='apply'
+            features = {
+                'cls_features':   pick('features'),
+                'features_proj':  pick('features_proj'),
+                'similarities':   pick('similarities'),  # pre-margin cosines
+                'cos_post':       pick('cos_post'),      # post-margin cosines (RoseFace)
+                'labels':         pick('labels'),
+                'logits':         pick('logits'),
+            }
+        return features
+
+
+    def analyze_feature_collapse(self, features):
+        print("\n=== FEATURE COLLAPSE ANALYSIS ===")
+        cls_features = features['cls_features'].numpy()
+        unique_patterns = self._count_unique_patterns(cls_features)
+        print(f"Estimated unique patterns: {unique_patterns}/100 classes")
+
+        feature_std = cls_features.std(axis=0).mean()
+        print(f"Average feature std: {feature_std:.4f}")
+
+        labels = features['labels'].numpy()
+        sample_size = min(1000, len(labels))
+        indices = np.random.choice(len(labels), sample_size, replace=False)
+        silhouette = silhouette_score(cls_features[indices], labels[indices])
+        print(f"Silhouette score: {silhouette:.3f}")
+
+        # centroid proximity count
+        class_features = {}
+        for i in range(100):
+            mask = labels == i
+            if mask.sum() > 0:
+                class_features[i] = cls_features[mask].mean(axis=0)
+        if class_features:
+            centroids = np.stack(list(class_features.values()))
+            d = np.linalg.norm(centroids[:, None] - centroids[None, :], axis=2)
+            thr = np.percentile(d[d > 0], 20)
+            close_pairs = (d < thr) & (d > 0)
+            classes_with_close_neighbors = close_pairs.sum(axis=1)
+            print(f"Classes with very similar features: {(classes_with_close_neighbors > 2).sum()}/100")
+
+        return {'unique_patterns': unique_patterns, 'feature_std': feature_std, 'silhouette': silhouette}
+
+    def analyze_geometric_patterns(self, features):
+        print("\n=== GEOMETRIC PATTERN ANALYSIS ===")
+        sims = features['similarities']  # pre-margin cosines [N,C]
+        print(f"Average max cosine: {sims.max(dim=1)[0].mean():.3f}")
+        print(f"Average min cosine: {sims.min(dim=1)[0].mean():.3f}")
+        print(f"Cosine std: {sims.std():.3f}")
+
+        high_sim_threshold = sims.mean() + sims.std()
+        high_sim_count = (sims > high_sim_threshold).sum(dim=1).float().mean()
+        print(f"Avg classes above (mean+std): {high_sim_count:.1f}/100")
+
+        labels = features['labels']
+        correct = sims.gather(1, labels.view(-1,1)).squeeze(1).mean().item()
+        wrong = (sims.sum(dim=1) - sims.gather(1, labels.view(-1,1)).squeeze(1)) / (sims.size(1)-1)
+        margin = (correct - wrong.mean().item())
+        print(f"Avg cosine margin (correct - mean wrong): {margin:.3f}")
+
+        # RoseFace-specific: if post cosines are present, compare deltas
+        if 'cos_post' in features and features['cos_post'].numel() > 0:
+            cos_post = features['cos_post']
+            # shift on target column
+            pos_pre = sims.gather(1, labels.view(-1,1))
+            pos_post = cos_post.gather(1, labels.view(-1,1))
+            shift = (pos_post - pos_pre).mean().item()
+            print(f"Avg target cosine shift (post - pre): {shift:.3f}")
+
+        return {
+            'max_cos': sims.max(dim=1)[0].mean().item(),
+            'cos_std': sims.std().item(),
+            'high_sim_classes': high_sim_count.item(),
+            'discrimination_margin': margin
+        }
+
+    def test_linear_probe(self, features, num_epochs=50):
+        print("\n=== LINEAR PROBE TEST ===")
+        X = features['cls_features']
+        y = features['labels']
+        n_train = int(0.8 * len(y))
+        X_train, y_train = X[:n_train], y[:n_train]
+        X_test, y_test = X[n_train:], y[n_train:]
+
+        probe = torch.nn.Linear(X_train.shape[1], 100).to(self.device)
+        opt = torch.optim.Adam(probe.parameters(), lr=0.01)
+
+        X_train = X_train.to(self.device); y_train = y_train.to(self.device)
+        X_test  = X_test.to(self.device);  y_test  = y_test.to(self.device)
+
+        best = 0.0
+        for epoch in range(num_epochs):
+            logits = probe(X_train)
+            loss = F.cross_entropy(logits, y_train)
+            opt.zero_grad(); loss.backward(); opt.step()
+            if epoch % 10 == 0:
+                with torch.no_grad():
+                    acc = (probe(X_test).argmax(dim=1) == y_test).float().mean().item()
+                    best = max(best, acc)
+                    print(f"  Epoch {epoch}: Test acc = {acc*100:.1f}%")
+        with torch.no_grad():
+            final = (probe(X_test).argmax(dim=1) == y_test).float().mean().item()
+            best = max(best, final)
+        print(f"Best linear probe accuracy: {best*100:.1f}%")
+        return best
+
+    def visualize_features(self, features, method='tsne', n_samples=2000):
+        print(f"\n=== FEATURE VISUALIZATION ({method.upper()}) ===")
+        cls_features = features['cls_features'].numpy()
+        labels = features['labels'].numpy()
+        n_samples = min(n_samples, len(labels))
+        idx = np.random.choice(len(labels), n_samples, replace=False)
+        X = cls_features[idx]; y = labels[idx]
+
+        print(f"Reducing {n_samples} samples to 2D...")
+        reducer = TSNE(n_components=2, random_state=42, perplexity=30) if method=='tsne' else PCA(n_components=2)
+        X2 = reducer.fit_transform(X)
+
+        plt.figure(figsize=(12,9))
+        scatter = plt.scatter(X2[:,0], X2[:,1], c=y, cmap='nipy_spectral', alpha=0.6, s=15)
+        plt.title(f'Feature Space Visualization ({method.upper()})'); plt.xlabel('Comp 1'); plt.ylabel('Comp 2')
+
+        print("Estimating visual clusters...")
+        silhouette_scores, K = [], list(range(30, 60, 5))
+        for k in K:
+            kmeans = KMeans(n_clusters=k, random_state=42, n_init=3)
+            cls_lbl = kmeans.fit_predict(X2)
+            silhouette_scores.append(silhouette_score(X2, cls_lbl))
+        best_k = K[int(np.argmax(silhouette_scores))]
+        kmeans = KMeans(n_clusters=best_k, random_state=42, n_init=5)
+        cluster_labels = kmeans.fit_predict(X2)
+        n_populated = len(np.unique(cluster_labels))
+        plt.text(0.02, 0.98, f'Estimated clusters: {n_populated}', transform=plt.gca().transAxes,
+                 va='top', fontsize=12, bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
+        cbar = plt.colorbar(scatter, ticks=np.arange(0,100,10)); cbar.set_label('Class', rotation=270, labelpad=15)
+        plt.tight_layout(); plt.show()
+        return X2, n_populated
+
+    def analyze_pentachora_usage(self):
+        print("\n=== PENTACHORA USAGE ANALYSIS ===")
+        print(f"Classes: {self.model.num_classes}")
+        print(f"Embed dim: {self.model.embed_dim}  |  Penta dim: {self.model.pentachora_dim}")
+        print(f"Head: {getattr(self.model,'head_type','legacy')}  |  Prototype: {getattr(self.model,'prototype_mode','n/a')}  |  Margin: {getattr(self.model,'margin_type','n/a')}")
+        if hasattr(self.model, 'to_pentachora_dim'):
+            if isinstance(self.model.to_pentachora_dim, torch.nn.Linear):
+                print(f"Projection: Linear {self.model.embed_dim}→{self.model.pentachora_dim}")
+            else:
+                print("Projection: Identity")
+
+        # Inter-class centroid similarity (legacy view)
+        centroids = self.model.get_class_centroids()
+        sim = centroids @ centroids.t()
+        mask = ~torch.eye(self.model.num_classes, dtype=bool, device=sim.device)
+        off = sim[mask]
+        print(f"\nCentroid sims: mean={off.mean():.3f}  max={off.max():.3f}  min={off.min():.3f}")
+
+        # Principal-angle overlap
+        mean_fro, std_fro = principal_angle_overlap(self.model.class_pentachora)
+        print(f"Principal-angle Fro overlap: mean={mean_fro:.3f} ± {std_fro:.3f} (lower is better)")
+
+        return {'mean_similarity': off.mean().item(), 'max_similarity': off.max().item(), 'mean_fro_overlap': mean_fro}
+
+    def run_full_analysis(self):
+        print("="*60); print("COMPREHENSIVE FEATURE ANALYSIS"); print("="*60)
+        print("\nExtracting features (margin_mode =", self.margin_mode, ") ...")
+        feats = self.extract_all_features(max_batches=50)
+        print(f"✓ Extracted features from {len(feats['labels'])} samples")
+
+        res = {}
+        res['collapse']  = self.analyze_feature_collapse(feats)
+        res['geometric'] = self.analyze_geometric_patterns(feats)
+
+        # Margin stats from PRE-margin cosines
+        mu, sig = margin_stats(feats['similarities'], feats['labels'])
+        print(f"PRE-margin Δ (pos - bestneg): mean={mu:.3f}, std={sig:.3f}")
+
+        # Face-usage heatmap
+        if 'features_proj' in feats and feats['features_proj'].numel() > 0:
+            heat = face_usage_heatmap(self.model, feats['features_proj'].to(self.device), feats['labels'].to(self.device), norm_type=getattr(self.model,'norm_type','l1'))
+            print("Face-usage heatmap computed [C,10] (display top 3 classes by mass):")
+            class_mass = heat.sum(dim=1)
+            top3 = torch.topk(class_mass, k=min(3, heat.size(0))).indices.tolist()
+            for c in top3:
+                print(f"  class {c}: {heat[c].cpu().numpy().round(3)}")
+
+        res['linear_probe'] = self.test_linear_probe(feats)
+        res['pentachora']   = self.analyze_pentachora_usage()
+
+        # Visualizations
+        _, n_tsne = self.visualize_features(feats, 'tsne')
+        _, n_pca  = self.visualize_features(feats, 'pca')
+        res['visual_clusters'] = {'tsne': n_tsne, 'pca': n_pca}
+
+        # Summary
+        print("\n" + "="*60); print("DIAGNOSIS SUMMARY"); print("="*60)
+        up = res['collapse']['unique_patterns']; lp = res['linear_probe']
+        if up <= 45 and lp <= 0.42:
+            print(f"🔴 Compact regime: {up} unique patterns; linear probe {lp*100:.1f}%")
+        elif up > 60:
+            print(f"✅ Diverse regime: {up} unique patterns; linear probe {lp*100:.1f}%")
+        else:
+            print(f"⚡ Partial bottleneck: {up} unique patterns; linear probe {lp*100:.1f}%")
+
+        return res
+
+    # ------------------------------
+    # Helpers (unchanged interface)
+    # ------------------------------
+    def _count_unique_patterns(self, features, method='elbow'):
+        X = features[:min(3000, len(features))]
+        if method == 'elbow':
+            inertias, K = [], list(range(20, 80, 5))
+            for k in K:
+                km = KMeans(n_clusters=k, random_state=42, n_init=3)
+                km.fit(X); inertias.append(km.inertia_)
+            diffs = np.diff(inertias); diffs2 = np.diff(diffs)
+            if len(diffs2) > 0:
+                elbow_idx = int(np.argmax(np.abs(diffs2))) + 1
+                est = K[elbow_idx]
+            else:
+                est = 41
+        else:
+            scores, K = [], list(range(30, 60, 5))
+            for k in K:
+                km = KMeans(n_clusters=k, random_state=42, n_init=3)
+                lbl = km.fit_predict(X)
+                scores.append(silhouette_score(X, lbl))
+            est = K[int(np.argmax(scores))]
+        return est
+
+# ============================================
+# QUICK TEST (RoseFace-aware)
+# ============================================
+
+def quick_41_percent_test(model, test_loader, device=None, apply_margin_eval=False):
+    """
+    If apply_margin_eval=True, pass targets to model at eval (margin applied).
+    Otherwise, evaluate without margin (classic).
+    """
+    print("="*60); print("41% ACCURACY CAP HYPOTHESIS TEST"); print("="*60)
+    model.eval()
+    device = device or next(model.parameters()).device
+
+    # 1) Accuracy
+    print("\n1. Verifying model accuracy...")
+    correct, total = 0, 0
+    with torch.no_grad():
+        for images, labels in tqdm(test_loader, desc="Testing"):
+            images = images.to(device)
+            labels = labels.to(device)
+            outputs = model(images, targets=labels) if apply_margin_eval else model(images)
+            pred = outputs['logits'].argmax(dim=1)
+            correct += (pred == labels).sum().item()
+            total += labels.size(0)
+    acc = 100 * correct / total
+    policy = "WITH margin" if apply_margin_eval else "NO margin"
+    print(f"   Test Accuracy ({policy}): {acc:.2f}%")
+
+    is_at_cap = abs(acc - 41.0) < 3.0
+    # 2) Focused analysis (small sample)
+    print("\n2. Analyzing feature patterns...")
+    margin_mode = 'apply' if apply_margin_eval else 'none'
+    analyzer = FeatureAnalyzer(model, test_loader, device=device, margin_mode=margin_mode)
+    feats = analyzer.extract_all_features(max_batches=20)
+    acc_rose5 = offline_head_eval_rose5(model, feats['features_proj'].to(device), feats['labels'])
+    print(f"Offline prototype eval (rose5, no margin): {acc_rose5*100:.2f}%")
+    collapse = analyzer.analyze_feature_collapse(feats)
+    pent = analyzer.analyze_pentachora_usage()
+
+    print("\n" + "="*60); print("VERDICT"); print("="*60)
+    if is_at_cap and collapse['unique_patterns'] <= 45:
+        print("🔴 41% CAP CONFIRMED")
+        print(f"   Acc: {acc:.1f}%  |  Unique patterns: {collapse['unique_patterns']}")
+        print("   Likely geometric bottleneck.")
+    elif collapse['unique_patterns'] <= 45:
+        print("⚠️ FEATURE BOTTLENECK DETECTED")
+        print(f"   {collapse['unique_patterns']} patterns; Acc={acc:.1f}%")
+    else:
+        print("✅ NO 41% BOTTLENECK")
+        print(f"   {collapse['unique_patterns']} patterns; Acc={acc:.1f}%")
+
+    return {
+        'accuracy': acc,
+        'unique_patterns': collapse['unique_patterns'],
+        'is_bottlenecked': collapse['unique_patterns'] <= 45,
+        'pentachora_similarity': pent['mean_similarity']
+    }
+
+@torch.no_grad()
+def offline_head_eval_rose5(model, features_proj, labels):
+    # compute z_l2 (dual-norm bridge)
+    z = features_proj
+    z_l2 = z / (z.norm(p=2, dim=-1, keepdim=True) + 1e-12)
+    # build rose5 prototypes [C,D]
+    V = torch.stack([p.vertices for p in model.class_pentachora], dim=0).to(z.device, z.dtype)  # [C,5,D]
+    V = V / (V.norm(p=2, dim=-1, keepdim=True) + 1e-12)
+    W = model.rose_face_weights.to(z.device, z.dtype)  # [10,5]
+    faces = torch.einsum('tf,cfd->ctd', W, V)
+    proto = (V.mean(dim=1) + 0.5 * faces.mean(dim=1))
+    proto = proto / (proto.norm(p=2, dim=-1, keepdim=True) + 1e-12)  # [C,D]
+    cos = z_l2 @ proto.t()                                           # [N,C]
+    acc = (cos.argmax(dim=1) == labels.to(z.device)).float().mean().item()
+    return acc
+
+# ==========================
+# RUN ANALYSIS (example)
+# ==========================
+
+if __name__ == "__main__":
+    print("Starting RoseFace-aware feature analysis...")
+    print(f"Model device: {next(model.parameters()).device}")
+    # Dataloaders
+    train_loader, test_loader, train_transforms = get_cifar100_dataloaders(batch_size=128)
+
+    # Quick test in BOTH modes (optional): compare accuracy
+    print("\nRunning quick 41% hypothesis test (NO margin at eval)...")
+    res_nom = quick_41_percent_test(model, test_loader, apply_margin_eval=False)
+
+    print("\nRunning quick 41% hypothesis test (WITH margin at eval)...")
+    res_mar = quick_41_percent_test(model, test_loader, apply_margin_eval=True)
+
+    # Full analysis with richer diagnostics (no margin at eval is typical)
+    analyzer = FeatureAnalyzer(model, test_loader, margin_mode='none')
+    full_results = analyzer.run_full_analysis()