Create 111m_prototype_v4_200x24_svd_eigh_kl_div.py

111m_prototype_v4_200x24_svd_eigh_kl_div.py

@@ -0,0 +1,436 @@
+"""
+SVAE - V=1024, D=24 (Validated Binding Constant)
+==================================================
+V=1024, D=24 -> CV=0.2916 (from sweep, confirmed)
+
+Deep encoder/decoder for 1024x24 = 24,576 matrix.
+Light KL on spectral shape (don't constrain magnitude).
+Row CV should be ~0.29 by dimensional law.
+
+pip install "git+https://github.com/AbstractEyes/geolip-core.git"
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+import torchvision.transforms as T
+import math
+
+try:
+    from geolip_core.linalg import svd as geolip_svd
+    from geolip_core.linalg.eigh import FLEigh
+    HAS_GEOLIP = True
+    print("Using geolip-core FLEigh (Faddeev-LeVerrier pipeline)")
+except ImportError:
+    HAS_GEOLIP = False
+    print("geolip-core not found, fallback to torch.svd_lowrank")
+
+
+# -- CM monitoring --
+
+def cayley_menger_vol2(points):
+    B, N, D = points.shape
+    gram = torch.bmm(points, points.transpose(1, 2))
+    norms = torch.diagonal(gram, dim1=1, dim2=2)
+    d2 = F.relu(norms.unsqueeze(2) + norms.unsqueeze(1) - 2 * gram)
+    cm = torch.zeros(B, N + 1, N + 1, device=points.device, dtype=points.dtype)
+    cm[:, 0, 1:] = 1.0
+    cm[:, 1:, 0] = 1.0
+    cm[:, 1:, 1:] = d2
+    k = N - 1
+    sign = (-1.0) ** (k + 1)
+    fact = math.factorial(k)
+    return sign * torch.linalg.det(cm.float()).to(points.dtype) / ((2 ** k) * (fact ** 2))
+
+
+def cv_of(emb, n_samples=200):
+    if emb.dim() != 2 or emb.shape[0] < 5:
+        return 0.0
+    N, D = emb.shape
+    pool = min(N, 512)
+    indices = torch.stack([torch.randperm(pool, device=emb.device)[:5] for _ in range(n_samples)])
+    vol2 = cayley_menger_vol2(emb[:pool][indices])
+    valid = vol2 > 1e-20
+    if valid.sum() < 10:
+        return 0.0
+    vols = vol2[valid].sqrt()
+    return (vols.std() / (vols.mean() + 1e-8)).item()
+
+
+def safe_gram_svd(M, use_fast=False):
+    """
+    use_fast=False: torch.svd_lowrank (randomized, always converges, slower)
+    use_fast=True:  Gram + eigh in fp64 (deterministic, fast, needs conditioning)
+    """
+    orig_dtype = M.dtype
+    if use_fast:
+        # Optimized: Gram + eigh in fp64
+        A = M.double()
+        G = torch.bmm(A.transpose(1, 2), A)
+        eigenvalues, V = torch.linalg.eigh(G)
+        eigenvalues = eigenvalues.flip(-1)
+        V = V.flip(-1)
+        S = torch.sqrt(eigenvalues.clamp(min=1e-24))
+        U = torch.bmm(A, V) / S.unsqueeze(1).clamp(min=1e-16)
+        Vh = V.transpose(-2, -1).contiguous()
+        return U.to(orig_dtype), S.to(orig_dtype), Vh.to(orig_dtype)
+    else:
+        # Stable: randomized, always converges
+        U, S, V = torch.svd_lowrank(M.float(), q=M.shape[-1], niter=4)
+        Vh = V.transpose(1, 2)
+        return U.to(orig_dtype), S.to(orig_dtype), Vh.to(orig_dtype)
+
+
+BINDING_CONSTANT = 0.29154
+
+
+# -- Data --
+
+def get_cifar10(batch_size=256):
+    transform = T.Compose([
+        T.ToTensor(),
+        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
+    ])
+    train_ds = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
+    test_ds = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
+    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=2)
+    test_loader = torch.utils.data.DataLoader(test_ds, batch_size=batch_size, shuffle=False, num_workers=2)
+    return train_loader, test_loader
+
+
+# -- SVAE --
+
+class SVAE(nn.Module):
+    def __init__(self, matrix_v=200, D=24):
+        super().__init__()
+        self.matrix_v = matrix_v
+        self.D = D
+        self.img_dim = 3 * 32 * 32
+        self.mat_dim = matrix_v * D  # 200*24 = 4,800
+
+        # Gradual expansion: 3072 -> 512 -> 1024 -> 4800 (max 4.7x per step)
+        self.encoder = nn.Sequential(
+            nn.Linear(self.img_dim, 512),
+            nn.GELU(),
+            nn.Linear(512, 1024),
+            nn.GELU(),
+            nn.Linear(1024, self.mat_dim),
+        )
+        self.decoder = nn.Sequential(
+            nn.Linear(self.mat_dim, 1024),
+            nn.GELU(),
+            nn.Linear(1024, 512),
+            nn.GELU(),
+            nn.Linear(512, self.img_dim),
+        )
+
+        # Spectral log-variance (shape regularization only)
+        self.logvar_head = nn.Sequential(
+            nn.Linear(1024, 128),   # tap from encoder hidden layer
+            nn.GELU(),
+            nn.Linear(128, D),
+        )
+        # Init logvar to small values so reparameterization starts gentle
+        nn.init.zeros_(self.logvar_head[-1].weight)
+        nn.init.constant_(self.logvar_head[-1].bias, -5.0)
+
+        # Prior: SHAPE only, not magnitude
+        # Normalized decay from 1.0 to ~0.14 in log space
+        # The prior says "S should decay smoothly" not "S should be small"
+        self.register_buffer('prior_log_mu', torch.linspace(0, -2, D))
+        self.register_buffer('prior_log_var', torch.ones(D))  # wide prior (var=e^1 ~2.7)
+
+        # Orthogonal init on last encoder layer for well-conditioned initial matrices
+        nn.init.orthogonal_(self.encoder[-1].weight)
+
+        self.use_fast_svd = False  # switched during training
+
+    def encode(self, images):
+        B = images.shape[0]
+        flat = images.reshape(B, -1)
+
+        # Run encoder with hidden tap for logvar
+        h1 = F.gelu(self.encoder[0](flat))      # 3072 -> 512
+        h2 = F.gelu(self.encoder[2](h1))         # 512 -> 1024
+        mat_flat = self.encoder[4](h2)            # 1024 -> mat_dim
+        M = mat_flat.reshape(B, self.matrix_v, self.D)
+
+        U, S, Vh = safe_gram_svd(M, use_fast=self.use_fast_svd)
+
+        # Log-variance from hidden (not full mat_dim - too expensive)
+        log_var = self.logvar_head(h2)
+
+        # Reparameterize on NORMALIZED spectrum (shape, not magnitude)
+        if self.training:
+            S_norm = S / (S[:, 0:1] + 1e-8)  # normalize by S[0]
+            log_S_norm = torch.log(S_norm.clamp(min=1e-8))
+            std = torch.exp(0.5 * log_var)
+            eps = torch.randn_like(std)
+            S_norm_sampled = torch.exp(log_S_norm + std * eps)
+            S_norm_sampled, _ = S_norm_sampled.sort(dim=-1, descending=True)
+            # Denormalize back
+            S_sampled = S_norm_sampled * S[:, 0:1]
+        else:
+            S_sampled = S
+            S_norm = S / (S[:, 0:1] + 1e-8)
+
+        return {
+            'U': U, 'S': S, 'S_sampled': S_sampled, 'Vt': Vh,
+            'S_norm': S / (S[:, 0:1] + 1e-8),
+            'M': M, 'log_var': log_var,
+        }
+
+    def decode_from_svd(self, U, S, Vt):
+        B = U.shape[0]
+        M_hat = torch.bmm(U * S.unsqueeze(1), Vt)
+        return self.decoder(M_hat.reshape(B, -1)).reshape(B, 3, 32, 32)
+
+    def spectral_kl(self, S_norm, log_var):
+        """KL on NORMALIZED spectrum shape. Magnitude-free."""
+        log_S = torch.log(S_norm.clamp(min=1e-8))
+        mu_q = log_S
+        var_q = torch.exp(log_var)
+        mu_p = self.prior_log_mu.unsqueeze(0)
+        var_p = torch.exp(self.prior_log_var).unsqueeze(0)
+        kl = 0.5 * (var_q / var_p + (mu_p - mu_q).pow(2) / var_p
+                     - 1 + torch.log(var_p / (var_q + 1e-8)))
+        return kl.sum(dim=-1).mean()
+
+    def forward(self, images):
+        svd = self.encode(images)
+        if self.use_fast_svd:
+            # Post-warmup: decode from sampled S, compute KL
+            recon = self.decode_from_svd(svd['U'], svd['S_sampled'], svd['Vt'])
+            kl = self.spectral_kl(svd['S_norm'], svd['log_var'])
+        else:
+            # Warmup: pure reconstruction, no KL noise
+            recon = self.decode_from_svd(svd['U'], svd['S'], svd['Vt'])
+            kl = torch.tensor(0.0, device=images.device)
+        return {'recon': recon, 'svd': svd, 'kl': kl}
+
+    @staticmethod
+    def effective_rank(S):
+        p = S / (S.sum(-1, keepdim=True) + 1e-8)
+        p = p.clamp(min=1e-8)
+        return (-(p * p.log()).sum(-1)).exp()
+
+
+# -- Training --
+
+def train(epochs=50, lr=1e-3, kl_weight=0.001, warmup_epochs=5, device='cuda'):
+    device = torch.device(device if torch.cuda.is_available() else 'cpu')
+    train_loader, test_loader = get_cifar10(batch_size=256)
+
+    model = SVAE(matrix_v=200, D=24).to(device)
+    opt = torch.optim.Adam(model.parameters(), lr=lr)
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"SVAE - V=200, D=24 (Validated: CV=0.2914)")
+    print(f"  Matrix: (200, 24) = 4,800 elements")
+    print(f"  Phase 1 (ep 1-{warmup_epochs}): randomized SVD, recon only")
+    print(f"  Phase 2 (ep {warmup_epochs+1}+): fp64 Gram+eigh, recon + KL(w={kl_weight})")
+    print(f"  Params: {total_params:,}")
+    print("=" * 95)
+    print(f"{'ep':>3} | {'loss':>7} {'recon':>7} {'kl':>7} {'t/ep':>5} | "
+          f"{'t_rec':>7} | "
+          f"{'S0':>7} {'SD':>6} {'ratio':>5} {'erank':>5} | "
+          f"{'row_cv':>7} {'svd':>5}")
+    print("-" * 95)
+
+    import time
+
+    for epoch in range(1, epochs + 1):
+        # SVD switch: randomized warmup -> optimized
+        if epoch == warmup_epochs + 1:
+            model.use_fast_svd = True
+            print(f"  >>> Switching to fp64 Gram+eigh SVD <<<")
+
+        model.train()
+        total_loss, total_recon, total_kl, n = 0, 0, 0, 0
+        t0 = time.time()
+
+        for images, labels in train_loader:
+            images = images.to(device)
+            opt.zero_grad()
+            out = model(images)
+
+            recon_loss = F.mse_loss(out['recon'], images)
+            kl = out['kl']
+            loss = recon_loss + kl_weight * kl
+            loss.backward()
+            opt.step()
+
+            total_loss += loss.item() * len(images)
+            total_recon += recon_loss.item() * len(images)
+            total_kl += kl.item() * len(images)
+            n += len(images)
+
+        sched.step()
+        epoch_time = time.time() - t0
+
+        if epoch % 2 == 0 or epoch <= 3 or epoch == warmup_epochs + 1:
+            model.eval()
+            test_recon, test_n = 0, 0
+            test_S, test_erank = None, 0
+            row_cvs = []
+            nb = 0
+
+            with torch.no_grad():
+                for images, labels in test_loader:
+                    images = images.to(device)
+                    out = model(images)
+                    test_recon += F.mse_loss(out['recon'], images).item() * len(images)
+                    test_n += len(images)
+                    test_erank += model.effective_rank(out['svd']['S']).mean().item()
+
+                    if nb < 3:
+                        for b in range(min(4, len(images))):
+                            row_cvs.append(cv_of(out['svd']['M'][b]))
+
+                    if test_S is None:
+                        test_S = out['svd']['S'].mean(0).cpu()
+                    else:
+                        test_S += out['svd']['S'].mean(0).cpu()
+                    nb += 1
+
+            test_erank /= nb
+            test_S /= nb
+            ratio = (test_S[0] / (test_S[-1] + 1e-8)).item()
+            mean_cv = sum(row_cvs) / len(row_cvs) if row_cvs else 0
+            svd_tag = "FAST" if model.use_fast_svd else "rand"
+
+            print(f"{epoch:3d} | {total_loss/n:7.4f} {total_recon/n:7.4f} "
+                  f"{total_kl/n:7.3f} {epoch_time:5.1f} | "
+                  f"{test_recon/test_n:7.4f} | "
+                  f"{test_S[0]:7.2f} {test_S[-1]:6.3f} {ratio:5.2f} "
+                  f"{test_erank:5.2f} | "
+                  f"{mean_cv:7.4f} {svd_tag:>5}")
+
+    # -- Final Analysis --
+    print()
+    print("=" * 90)
+    print("FINAL ANALYSIS")
+    print("=" * 90)
+
+    model.eval()
+    all_S, all_recon_err, all_labels = [], [], []
+    all_row_cvs = []
+
+    with torch.no_grad():
+        for images, labels in test_loader:
+            images = images.to(device)
+            out = model(images)
+            all_S.append(out['svd']['S'].cpu())
+            all_recon_err.append(
+                F.mse_loss(out['recon'], images, reduction='none')
+                .mean(dim=(1, 2, 3)).cpu())
+            all_labels.append(labels.cpu())
+            for b in range(min(4, len(images))):
+                all_row_cvs.append(cv_of(out['svd']['M'][b]))
+
+    all_S = torch.cat(all_S)
+    all_recon_err = torch.cat(all_recon_err)
+    all_labels = torch.cat(all_labels)
+    erank = model.effective_rank(all_S)
+    mean_cv = sum(all_row_cvs) / len(all_row_cvs)
+
+    print(f"\n  V=1024, D=24 (validated CV=0.2916)")
+    print(f"  Recon MSE: {all_recon_err.mean():.6f} +/- {all_recon_err.std():.6f}")
+    print(f"  Effective rank: {erank.mean():.2f} +/- {erank.std():.2f}")
+    print(f"  Row CV: {mean_cv:.4f} (target: {BINDING_CONSTANT}, delta: {abs(mean_cv - BINDING_CONSTANT):.4f})")
+
+    # Spectrum
+    S_mean = all_S.mean(0)
+    S_norm = S_mean / (S_mean[0] + 1e-8)
+    total_energy = (S_mean ** 2).sum()
+    print(f"\n  Singular value profile (raw and normalized):")
+    cumulative = 0
+    for i in range(len(S_mean)):
+        e = (S_mean[i] ** 2).item()
+        cumulative += e
+        pct = cumulative / total_energy * 100
+        bar = "#" * int(S_norm[i].item() * 30)
+        print(f"    S[{i:2d}]: {S_mean[i]:8.3f}  norm={S_norm[i]:.4f}  cum={pct:5.1f}%  {bar}")
+
+    # Per-class
+    cifar_names = ['plane', 'car', 'bird', 'cat', 'deer',
+                   'dog', 'frog', 'horse', 'ship', 'truck']
+    print(f"\n  Per-class:")
+    print(f"    {'cls':>6}  {'recon':>8}  {'erank':>6}  {'S0':>7}  {'SD':>7}  {'ratio':>6}")
+    for c in range(10):
+        mask = all_labels == c
+        rc = all_recon_err[mask].mean().item()
+        er = erank[mask].mean().item()
+        s0 = all_S[mask, 0].mean().item()
+        sd = all_S[mask, -1].mean().item()
+        r = s0 / (sd + 1e-8)
+        print(f"    {cifar_names[c]:>6}  {rc:8.6f}  {er:6.2f}  {s0:7.3f}  {sd:7.3f}  {r:6.2f}")
+
+    # -- Recon grid --
+    print(f"\n  Saving reconstruction grid...")
+    import matplotlib
+    matplotlib.use('Agg')
+    import matplotlib.pyplot as plt
+
+    mean_t = torch.tensor([0.4914, 0.4822, 0.4465]).reshape(1, 3, 1, 1).to(device)
+    std_t = torch.tensor([0.2470, 0.2435, 0.2616]).reshape(1, 3, 1, 1).to(device)
+
+    model.eval()
+    with torch.no_grad():
+        images, labels = next(iter(test_loader))
+        images = images.to(device)
+        out = model(images)
+
+        selected_idx = []
+        for c in range(10):
+            class_idx = (labels == c).nonzero(as_tuple=True)[0]
+            selected_idx.extend(class_idx[:2].tolist())
+
+        orig = images[selected_idx]
+        U = out['svd']['U'][selected_idx]
+        S = out['svd']['S'][selected_idx]
+        Vt = out['svd']['Vt'][selected_idx]
+
+        mode_counts = [1, 4, 8, 16, 24]
+        prog_recons = []
+        for nm in mode_counts:
+            r = model.decode_from_svd(U[:, :, :nm], S[:, :nm], Vt[:, :nm, :])
+            prog_recons.append(r)
+
+    def denorm(t):
+        return (t * std_t + mean_t).clamp(0, 1).cpu()
+
+    n_samples = len(selected_idx)
+    n_cols = 2 + len(mode_counts)
+    fig, axes = plt.subplots(n_samples, n_cols, figsize=(n_cols * 1.5, n_samples * 1.5))
+    col_titles = ['Original'] + [f'{m} modes' for m in mode_counts] + ['|Err|x5']
+
+    for i in range(n_samples):
+        axes[i, 0].imshow(denorm(orig[i:i+1])[0].permute(1, 2, 0).numpy())
+        for j, r in enumerate(prog_recons):
+            axes[i, j+1].imshow(denorm(r[i:i+1])[0].permute(1, 2, 0).numpy())
+        err_col = 1 + len(prog_recons)
+        diff = (denorm(orig[i:i+1]) - denorm(prog_recons[-1][i:i+1])).abs() * 5
+        axes[i, err_col].imshow(diff.clamp(0, 1)[0].permute(1, 2, 0).numpy())
+        c = labels[selected_idx[i]].item()
+        axes[i, 0].set_ylabel(cifar_names[c], fontsize=8, rotation=0, labelpad=35)
+
+    for j, title in enumerate(col_titles):
+        axes[0, j].set_title(title, fontsize=8)
+    for ax in axes.flat:
+        ax.axis('off')
+
+    plt.tight_layout()
+    plt.savefig('/content/svae_recon_grid.png', dpi=200, bbox_inches='tight')
+    print(f"  Saved to /content/svae_recon_grid.png")
+    try:
+        plt.show()
+    except:
+        pass
+    plt.close()
+
+
+if __name__ == "__main__":
+    train()