""" SVAE — SVD Autoencoder with Geometric Attractors =================================================== A matrix-valued autoencoder where the latent space is a (V, D) matrix decomposed by SVD. Rows are normalized to S^(D-1), making the geometric structure architectural rather than loss-dependent. Two key mechanisms: 1. Sphere normalization: F.normalize(M, dim=-1) constrains rows to unit vectors on S^(D-1). This bounds the Gram matrix, eliminates training instabilities, and makes the CV a structural property of (V, D). 2. Soft hand: An oscillatory counterweight that boosts reconstruction gradients when geometry is near target, and penalizes CV drift when geometry is far from target. Provides positive momentum, not just penalty. Architecture: Image → MLP → M ∈ ℝ^(V×D) → normalize → SVD → MLP → Recon Repository: AbstractEyes/geolip-core """ import torch import torch.nn as nn import torch.nn.functional as F import torchvision import torchvision.transforms as T import math import time # ── SVD Backend ────────────────────────────────────────────────── try: from geolip_core.linalg.eigh import FLEigh, _FL_MAX_N HAS_FL = True except ImportError: HAS_FL = False def gram_eigh_svd_fp64(A): """Thin SVD via Gram matrix + eigh, computed entirely in fp64. fp64 is essential: Gram entries scale as S₀², and fp32 (~7 digits) causes catastrophic collapses when the condition number exceeds ~100. fp64 (~15 digits) eliminates this failure mode entirely. Args: A: (B, M, N) tensor, M >= N Returns: U (B,M,N), S (B,N), Vh (B,N,N) — singular values descending. """ orig_dtype = A.dtype with torch.amp.autocast('cuda', enabled=False): A_d = A.double() G = torch.bmm(A_d.transpose(1, 2), A_d) 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_d, 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) def svd_fp64(A): """Auto-dispatch SVD with fp64 internals. N <= 12 + FLEigh available: Gram in fp64, FL eigh (compilable). N > 12 or CPU: Gram + torch.linalg.eigh in fp64. Triton bypassed — fp32-only hardware, incompatible with fp64. """ B, M, N = A.shape if HAS_FL and N <= _FL_MAX_N and A.is_cuda: orig_dtype = A.dtype with torch.amp.autocast('cuda', enabled=False): A_d = A.double() G = torch.bmm(A_d.transpose(1, 2), A_d) eigenvalues, V = FLEigh()(G.float()) # FL needs fp32 input eigenvalues = eigenvalues.double().flip(-1) V = V.double().flip(-1) S = torch.sqrt(eigenvalues.clamp(min=1e-24)) U = torch.bmm(A_d, 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: return gram_eigh_svd_fp64(A) # ── Cayley-Menger CV Monitoring ────────────────────────────────── def cayley_menger_vol2(points): """Squared simplex volume via Cayley-Menger determinant, in fp64. Args: points (B, N, D) — B simplices, each with N vertices in D dims. Returns: (B,) squared volumes. """ B, N, D = points.shape pts = points.double() gram = torch.bmm(pts, pts.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=torch.float64) 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) / ((2 ** k) * (fact ** 2)) def cv_of(emb, n_samples=200): """CV of pentachoron volumes for a single embedding matrix. Measures geometric regularity: low CV = regular, high CV = irregular. Args: emb (V, D) tensor. Returns: float CV value, or 0.0 if insufficient data. """ 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() # ── 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, 3 * 32 * 32, 10, \ ['plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'] def get_tiny_imagenet(batch_size=256): """TinyImageNet via HuggingFace: 200 classes, 64x64, 100K train / 10K val.""" from datasets import load_dataset ds = load_dataset('zh-plus/tiny-imagenet') mean = (0.4802, 0.4481, 0.3975) std = (0.2770, 0.2691, 0.2821) transform = T.Compose([ T.ToTensor(), T.Normalize(mean, std), ]) class HFImageDataset(torch.utils.data.Dataset): def __init__(self, hf_split, transform): self.data = hf_split self.transform = transform def __len__(self): return len(self.data) def __getitem__(self, idx): item = self.data[idx] img = item['image'] if img.mode != 'RGB': img = img.convert('RGB') return self.transform(img), item['label'] train_ds = HFImageDataset(ds['train'], transform) val_ds = HFImageDataset(ds['valid'], transform) train_loader = torch.utils.data.DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=2) val_loader = torch.utils.data.DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=2) class_names = [f'c{i:03d}' for i in range(200)] return train_loader, val_loader, 3 * 64 * 64, 200, class_names # ── SVAE Model ─────────────────────────────────────────────────── class SVAE(nn.Module): """SVD Autoencoder with sphere-normalized matrix latent space. The encoder produces a (V, D) matrix whose rows are normalized to S^(D-1). The SVD decomposes alignment structure (U, V) from spectral magnitudes (S). The decoder reconstructs from the full SVD: M̂ = UΣVᵀ. Args: matrix_v: Number of rows V (vocabulary size / overcomplete factor) D: Embedding dimension (number of singular values) img_dim: Flattened image dimension (3*H*W) hidden: Hidden layer width (auto-scaled if None) """ def __init__(self, matrix_v=96, D=24, img_dim=3072, hidden=None): super().__init__() self.matrix_v = matrix_v self.D = D self.img_dim = img_dim self.mat_dim = matrix_v * D h = hidden or max(512, min(2048, img_dim // 4)) self.encoder = nn.Sequential( nn.Linear(self.img_dim, h), nn.GELU(), nn.Linear(h, h), nn.GELU(), nn.Linear(h, self.mat_dim), ) self.decoder = nn.Sequential( nn.Linear(self.mat_dim, h), nn.GELU(), nn.Linear(h, h), nn.GELU(), nn.Linear(h, self.img_dim), ) nn.init.orthogonal_(self.encoder[-1].weight) def encode(self, images): B = images.shape[0] M = self.encoder(images.reshape(B, -1)).reshape(B, self.matrix_v, self.D) M = F.normalize(M, dim=-1) # rows to S^(D-1) U, S, Vh = svd_fp64(M) return {'U': U, 'S': S, 'Vt': Vh, 'M': M} def decode_from_svd(self, U, S, Vt): B = U.shape[0] M_hat = torch.bmm(U * S.unsqueeze(1), Vt) flat = self.decoder(M_hat.reshape(B, -1)) # Infer spatial dims from img_dim: 3*H*W hw = self.img_dim // 3 h = int(hw ** 0.5) return flat.reshape(B, 3, h, h) def forward(self, images): svd = self.encode(images) recon = self.decode_from_svd(svd['U'], svd['S'], svd['Vt']) return {'recon': recon, 'svd': svd} @staticmethod def effective_rank(S): """Shannon entropy effective rank of singular value spectrum.""" 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=100, lr=1e-3, V=256, D=24, target_cv=0.125, cv_weight=0.3, boost=0.5, sigma=0.15, dataset='cifar10', device='cuda'): """Train the SVAE with sphere normalization + soft hand. Args: epochs: Training epochs lr: Learning rate for Adam V: Matrix rows (vocabulary size) D: Embedding dimension target_cv: CV attractor target for soft hand cv_weight: Maximum CV penalty weight (far from target) boost: Maximum reconstruction boost factor (near target) sigma: Gaussian transition width for proximity dataset: 'cifar10' or 'tiny_imagenet' device: Training device """ device = torch.device(device if torch.cuda.is_available() else 'cpu') if dataset == 'tiny_imagenet': train_loader, test_loader, img_dim, n_classes, class_names = get_tiny_imagenet(batch_size=256) else: train_loader, test_loader, img_dim, n_classes, class_names = get_cifar10(batch_size=256) img_h = int((img_dim // 3) ** 0.5) model = SVAE(matrix_v=V, D=D, img_dim=img_dim).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()) # ── Header ── svd_backend = f"fp64 Gram+eigh (FL={'available, N<=12' if HAS_FL else 'not available'})" print(f"Using geolip-core SVD ({svd_backend})") print(f"SVAE - V={V}, D={D}, rows on S^{D-1} + soft hand") print(f" Dataset: {dataset} ({img_h}x{img_h}, {n_classes} classes)") print(f" Matrix: ({V}, {D}) = {V*D} elements, rows normalized") print(f" SVD: fp64 Gram+eigh") print(f" Sphere: rows on S^{D-1} (structural geometry)") print(f" Soft hand: boost={1+boost:.1f}x near CV={target_cv}, penalty={cv_weight} far") print(f" Params: {total_params:,}") print("=" * 90) print(f" {'ep':>3} | {'loss':>7} {'recon':>7} {'t/ep':>5} | " f"{'t_rec':>7} | " f"{'S0':>6} {'SD':>6} {'ratio':>5} {'erank':>5} | " f"{'row_cv':>7} {'prox':>5} {'rw':>5}") print("-" * 90) # ── Training loop ── for epoch in range(1, epochs + 1): model.train() total_loss, total_recon, n = 0, 0, 0 last_cv = target_cv last_prox = 1.0 recon_w = 1.0 + boost t0 = time.time() for batch_idx, (images, labels) in enumerate(train_loader): images = images.to(device) opt.zero_grad() out = model(images) recon_loss = F.mse_loss(out['recon'], images) # Measure CV and compute proximity (every 10th batch) with torch.no_grad(): if batch_idx % 10 == 0: current_cv = cv_of(out['svd']['M'][0]) if current_cv > 0: last_cv = current_cv delta = last_cv - target_cv last_prox = math.exp(-delta**2 / (2 * sigma**2)) # Soft hand: boost recon near target, penalize CV far from target recon_w = 1.0 + boost * last_prox cv_pen = cv_weight * (1.0 - last_prox) cv_l = (last_cv - target_cv) ** 2 loss = recon_w * recon_loss + cv_pen * cv_l loss.backward() opt.step() total_loss += loss.item() * len(images) total_recon += recon_loss.item() * len(images) n += len(images) sched.step() epoch_time = time.time() - t0 # ── Evaluation (every 2 epochs + first 3) ── if epoch % 2 == 0 or epoch <= 3: 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 < 5: 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 print(f" {epoch:3d} | {total_loss/n:7.4f} {total_recon/n:7.4f} {epoch_time:5.1f} | " f"{test_recon/test_n:7.4f} | " f"{test_S[0]:6.3f} {test_S[-1]:6.3f} {ratio:5.2f} " f"{test_erank:5.2f} | " f"{mean_cv:7.4f} {last_prox:5.3f} {recon_w:5.2f}") # ── Final Analysis ── print() print("=" * 85) print("FINAL ANALYSIS") print("=" * 85) 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(8, 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={V}, D={D}, rows on S^{D-1}") print(f" Target CV: {target_cv}") 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}") S_mean = all_S.mean(0) total_energy = (S_mean ** 2).sum() print(f"\n Singular value profile:") 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_mean[i].item() * 30 / (S_mean[0].item() + 1e-8)) print(f" S[{i:2d}]: {S_mean[i]:8.4f} cum={pct:5.1f}% {bar}") # Per-class summary (cap at 20 classes for readability) show_classes = min(n_classes, 20) print(f"\n Per-class (showing {show_classes}/{n_classes}):") print(f" {'cls':>8} {'recon':>8} {'erank':>6} {'S0':>7} {'SD':>7} {'ratio':>6}") for c in range(show_classes): mask = all_labels == c if mask.sum() == 0: continue 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) name = class_names[c] if c < len(class_names) else f'cls_{c}' print(f" {name:>8} {rc:8.6f} {er:6.2f} {s0:7.4f} {sd:7.4f} {r:6.2f}") # ── Reconstruction Grid ── print(f"\n Saving reconstruction grid...") import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt if dataset == 'tiny_imagenet': mean_t = torch.tensor([0.4802, 0.4481, 0.3975]).reshape(1, 3, 1, 1).to(device) std_t = torch.tensor([0.2770, 0.2691, 0.2821]).reshape(1, 3, 1, 1).to(device) else: 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) # Select up to 2 samples from up to 10 classes grid_classes = min(n_classes, 10) 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(grid_classes): class_idx = (labels == c).nonzero(as_tuple=True)[0] selected_idx.extend(class_idx[:2].tolist()) if not selected_idx: selected_idx = list(range(min(20, len(images)))) 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, D] 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() name = class_names[c] if c < len(class_names) else f'{c}' axes[i, 0].set_ylabel(name, fontsize=7, rotation=0, labelpad=40) 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(epochs=200, V=256, D=24, target_cv=0.45, dataset='tiny_imagenet')