Create constellation_vs_rope.py

constellation_vs_rope.py

@@ -0,0 +1,477 @@
+#!/usr/bin/env python3
+"""
+RoPE Attention vs Constellation Relay
+========================================
+Two RoPE variants:
+  1. Standard RoPE (Su et al.) — fixed base frequency 10000
+  2. NTK-aware RoPE — scaled base frequency for longer contexts
+
+Same battery of tests: single pass, depth stability, interleaved.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import math
+import time
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+torch.manual_seed(42)
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+HAS_FP8 = hasattr(torch, 'float8_e4m3fn')
+
+
+def compute_cv(points, n_samples=2000, n_points=5):
+    N = points.shape[0]
+    if N < n_points: return float('nan')
+    points = F.normalize(points.to(DEVICE).float(), dim=-1)
+    vols = []
+    for _ in range(n_samples):
+        idx = torch.randperm(min(N, 10000), device=DEVICE)[:n_points]
+        pts = points[idx].unsqueeze(0)
+        gram = torch.bmm(pts, pts.transpose(1, 2))
+        norms = torch.diagonal(gram, dim1=1, dim2=2)
+        d2 = norms.unsqueeze(2) + norms.unsqueeze(1) - 2 * gram
+        d2 = F.relu(d2)
+        cm = torch.zeros(1, 6, 6, device=DEVICE, dtype=torch.float32)
+        cm[:, 0, 1:] = 1; cm[:, 1:, 0] = 1; cm[:, 1:, 1:] = d2
+        v2 = -torch.linalg.det(cm) / 9216
+        if v2[0].item() > 1e-20:
+            vols.append(v2[0].sqrt().cpu())
+    if len(vols) < 50: return float('nan')
+    vt = torch.stack(vols)
+    return (vt.std() / (vt.mean() + 1e-8)).item()
+
+
+def eff_dim(x):
+    x_c = x - x.mean(0, keepdim=True)
+    _, S, _ = torch.linalg.svd(x_c[:512].float(), full_matrices=False)
+    p = S / S.sum()
+    return p.pow(2).sum().reciprocal().item()
+
+
+def uniform_sphere(n, d):
+    return F.normalize(torch.randn(n, d), dim=-1)
+
+
+# ══════════════════════════════════════════════════════════════════
+# RoPE IMPLEMENTATIONS
+# ══════════════════════════════════════════════════════════════════
+
+class RotaryEmbedding(nn.Module):
+    """Standard RoPE — fixed sinusoidal rotation frequencies."""
+    def __init__(self, dim, base=10000.0):
+        super().__init__()
+        self.dim = dim
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+
+    def forward(self, seq_len, device):
+        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
+        freqs = torch.einsum('i,j->ij', t, self.inv_freq)  # (S, dim/2)
+        emb = torch.cat([freqs, freqs], dim=-1)  # (S, dim)
+        return emb.cos(), emb.sin()
+
+
+class NTKRotaryEmbedding(nn.Module):
+    """NTK-aware RoPE — scaled base for extended context."""
+    def __init__(self, dim, base=10000.0, scale_factor=4.0):
+        super().__init__()
+        self.dim = dim
+        # NTK scaling: base^(dim/(dim-2)) * scale_factor
+        scaled_base = base * (scale_factor ** (dim / (dim - 2)))
+        inv_freq = 1.0 / (scaled_base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+
+    def forward(self, seq_len, device):
+        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
+        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+        emb = torch.cat([freqs, freqs], dim=-1)
+        return emb.cos(), emb.sin()
+
+
+def apply_rotary(x, cos, sin):
+    """Apply rotary embeddings to Q or K: (B, H, S, d)."""
+    d = x.shape[-1]
+    x1 = x[..., :d//2]
+    x2 = x[..., d//2:]
+    cos = cos[:x.shape[-2], :d//2].unsqueeze(0).unsqueeze(0)  # (1, 1, S, d/2)
+    sin = sin[:x.shape[-2], :d//2].unsqueeze(0).unsqueeze(0)
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+
+# ══════════════════════════════════════════════════════════════════
+# ATTENTION BLOCKS
+# ══════════════════════════════════════════════════════════════════
+
+class VanillaAttnBlock(nn.Module):
+    """Standard self-attention — no position encoding."""
+    def __init__(self, dim, n_heads=4):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = dim // n_heads
+        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.out_proj = nn.Linear(dim, dim, bias=False)
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x):
+        B, S, D = x.shape
+        x_n = self.norm(x)
+        qkv = self.qkv(x_n).reshape(B, S, 3, self.n_heads, self.head_dim)
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        attn = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5)
+        attn = attn.softmax(dim=-1)
+        out = (attn @ v).transpose(1, 2).reshape(B, S, D)
+        return x + self.out_proj(out)
+
+
+class RoPEAttnBlock(nn.Module):
+    """Self-attention with Rotary Position Embeddings."""
+    def __init__(self, dim, n_heads=4, rope_type='standard', rope_base=10000.0,
+                 ntk_scale=4.0):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = dim // n_heads
+        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.out_proj = nn.Linear(dim, dim, bias=False)
+        self.norm = nn.LayerNorm(dim)
+
+        if rope_type == 'standard':
+            self.rope = RotaryEmbedding(self.head_dim, base=rope_base)
+        elif rope_type == 'ntk':
+            self.rope = NTKRotaryEmbedding(self.head_dim, base=rope_base,
+                                            scale_factor=ntk_scale)
+        self.rope_type = rope_type
+
+    def forward(self, x):
+        B, S, D = x.shape
+        x_n = self.norm(x)
+        qkv = self.qkv(x_n).reshape(B, S, 3, self.n_heads, self.head_dim)
+        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, B, H, S, hd)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        # Apply RoPE to Q and K
+        cos, sin = self.rope(S, x.device)
+        q = apply_rotary(q, cos, sin)
+        k = apply_rotary(k, cos, sin)
+
+        attn = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5)
+        attn = attn.softmax(dim=-1)
+        out = (attn @ v).transpose(1, 2).reshape(B, S, D)
+        return x + self.out_proj(out)
+
+
+# ══════════════════════════════════════════════════════════════════
+# CONSTELLATION RELAY (copy from v2)
+# ══════════════════════════════════════════════════════════════════
+
+class ConstellationRelay(nn.Module):
+    def __init__(self, input_dim, patch_dim=16, n_anchors=16, n_phases=3,
+                 pw_hidden=32, gate_init=-3.0):
+        super().__init__()
+        assert input_dim % patch_dim == 0
+        self.input_dim = input_dim
+        self.patch_dim = patch_dim
+        self.n_patches = input_dim // patch_dim
+        self.n_anchors = n_anchors
+        self.n_phases = n_phases
+        P, A, d = self.n_patches, n_anchors, patch_dim
+
+        home = torch.empty(P, A, d)
+        nn.init.xavier_normal_(home.view(P * A, d))
+        home = F.normalize(home.view(P, A, d), dim=-1)
+        self.register_buffer('home', home)
+        self.anchors = nn.Parameter(home.clone())
+
+        tri_dim = n_phases * A
+        self.pw_w1 = nn.Parameter(torch.empty(P, tri_dim, pw_hidden))
+        self.pw_b1 = nn.Parameter(torch.zeros(1, P, pw_hidden))
+        self.pw_w2 = nn.Parameter(torch.empty(P, pw_hidden, d))
+        self.pw_b2 = nn.Parameter(torch.zeros(1, P, d))
+        for p in range(P):
+            nn.init.xavier_normal_(self.pw_w1.data[p])
+            nn.init.xavier_normal_(self.pw_w2.data[p])
+        self.pw_norm = nn.LayerNorm(d)
+        self.gates = nn.Parameter(torch.full((P,), gate_init))
+        self.norm = nn.LayerNorm(input_dim)
+
+    def drift(self):
+        h = F.normalize(self.home, dim=-1)
+        c = F.normalize(self.anchors, dim=-1)
+        cos = (h * c).sum(dim=-1).clamp(-1 + 1e-7, 1 - 1e-7)
+        return torch.acos(cos)
+
+    def at_phase(self, t):
+        h = F.normalize(self.home, dim=-1)
+        c = F.normalize(self.anchors, dim=-1)
+        omega = self.drift().unsqueeze(-1)
+        sin_omega = omega.sin().clamp(min=1e-7)
+        return (torch.sin((1 - t) * omega) / sin_omega * h +
+                torch.sin(t * omega) / sin_omega * c)
+
+    def forward(self, x):
+        B, D = x.shape
+        P, A, d = self.n_patches, self.n_anchors, self.patch_dim
+        x_n = self.norm(x)
+        patches = x_n.reshape(B, P, d)
+        patches_n = F.normalize(patches, dim=-1)
+
+        # Multi-phase triangulation
+        phases = torch.linspace(0, 1, self.n_phases).tolist()
+        tris = []
+        for t in phases:
+            anchors_t = F.normalize(self.at_phase(t), dim=-1)
+            cos = torch.einsum('bpd,pad->bpa', patches_n, anchors_t)
+            tris.append(1.0 - cos)
+        tri = torch.cat(tris, dim=-1)
+
+        # Patchwork
+        h = torch.einsum('bpt,pth->bph', tri, self.pw_w1) + self.pw_b1
+        h = F.gelu(h)
+        pw_out = torch.einsum('bph,phd->bpd', h, self.pw_w2) + self.pw_b2
+        pw_out = self.pw_norm(pw_out)
+
+        gate = self.gates.sigmoid().unsqueeze(0).unsqueeze(-1)
+        blended = gate * pw_out + (1 - gate) * patches
+        out = blended.reshape(B, D)
+        return x + out
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST SUITE
+# ══════════════════════════════════════════════════════════════════
+
+N = 2000
+D = 128
+N_CV = 2000
+
+print("=" * 90)
+print("RoPE ATTENTION vs CONSTELLATION RELAY")
+print(f"  Input dim: {D}, Sequence length: {N}")
+print(f"  Device: {DEVICE}")
+print("=" * 90)
+
+pts = uniform_sphere(N, D).to(DEVICE)
+cv_base = compute_cv(pts, N_CV)
+ed_base = eff_dim(pts)
+print(f"  Baseline: CV={cv_base:.4f}  eff_dim={ed_base:.1f}")
+
+# Build all architectures
+configs = {
+    'vanilla':    lambda: VanillaAttnBlock(D, 8).to(DEVICE),
+    'rope_std':   lambda: RoPEAttnBlock(D, 8, 'standard', 10000).to(DEVICE),
+    'rope_ntk':   lambda: RoPEAttnBlock(D, 8, 'ntk', 10000, 4.0).to(DEVICE),
+    'relay':      lambda: ConstellationRelay(D, 16, 16, 3, 32).to(DEVICE),
+}
+
+
+# ── TEST 1: Single pass comparison ──
+print(f"\n{'━'*90}")
+print("TEST 1: Single pass — all architectures")
+print(f"{'━'*90}")
+print(f"  {'arch':>12} {'params':>8} {'CV_out':>8} {'CV_norm':>8} "
+      f"{'cos_orig':>10} {'eff_dim':>8}")
+
+for name, builder in configs.items():
+    module = builder()
+    np_ = sum(p.numel() for p in module.parameters())
+    with torch.no_grad():
+        if name == 'relay':
+            out = module(pts)
+        else:
+            out = module(pts.unsqueeze(0)).squeeze(0)
+    cv = compute_cv(out, N_CV)
+    cv_n = compute_cv(F.normalize(out, dim=-1), N_CV)
+    cos = (F.normalize(pts, dim=-1) * F.normalize(out, dim=-1)).sum(-1).mean().item()
+    ed = eff_dim(out)
+    print(f"  {name:>12} {np_:>8,} {cv:>8.4f} {cv_n:>8.4f} {cos:>10.6f} {ed:>8.1f}")
+
+
+# ── TEST 2: Depth sweep — 16 layers each ──
+print(f"\n{'━'*90}")
+print("TEST 2: Depth sweep — 16 layers, all architectures")
+print(f"{'━'*90}")
+
+checkpoints = [1, 2, 4, 8, 12, 16]
+
+for name, builder in configs.items():
+    print(f"\n  {name}:")
+    print(f"  {'depth':>6} {'CV':>8} {'CV_n':>8} {'eff_d':>8} {'cos_orig':>10}")
+
+    stack = nn.ModuleList([builder() for _ in range(16)])
+    x = pts.clone()
+    for i, layer in enumerate(stack):
+        with torch.no_grad():
+            if name == 'relay':
+                x = layer(x)
+            else:
+                x = layer(x.unsqueeze(0)).squeeze(0)
+        if (i + 1) in checkpoints:
+            cv = compute_cv(x, N_CV)
+            cv_n = compute_cv(F.normalize(x, dim=-1), N_CV)
+            ed = eff_dim(x)
+            cos = (F.normalize(pts, dim=-1) * F.normalize(x, dim=-1)).sum(-1).mean().item()
+            print(f"  {i+1:>6} {cv:>8.4f} {cv_n:>8.4f} {ed:>8.1f} {cos:>10.6f}")
+
+
+# ── TEST 3: Interleaved — each attention type + relay ──
+print(f"\n{'━'*90}")
+print("TEST 3: Interleaved — [attn type] → relay → [attn type] → relay → ...")
+print(f"{'━'*90}")
+
+for attn_name in ['vanilla', 'rope_std', 'rope_ntk']:
+    print(f"\n  {attn_name} + relay interleaved:")
+    print(f"  {'step':>6} {'type':>8} {'CV_n':>8} {'eff_d':>8} {'cos_orig':>10}")
+
+    attn_builder = configs[attn_name]
+    attn_layers = nn.ModuleList([attn_builder() for _ in range(8)])
+    relay_layers = nn.ModuleList([
+        ConstellationRelay(D, 16, 16, 3, 32).to(DEVICE) for _ in range(8)])
+
+    x = pts.clone()
+    step = 0
+    for i in range(8):
+        # Attention step
+        with torch.no_grad():
+            x = attn_layers[i](x.unsqueeze(0)).squeeze(0)
+        step += 1
+        if step in checkpoints:
+            cv_n = compute_cv(F.normalize(x, dim=-1), N_CV)
+            ed = eff_dim(x)
+            cos = (F.normalize(pts, dim=-1) * F.normalize(x, dim=-1)).sum(-1).mean().item()
+            print(f"  {step:>6} {'attn':>8} {cv_n:>8.4f} {ed:>8.1f} {cos:>10.6f}")
+
+        # Relay step
+        with torch.no_grad():
+            x = relay_layers[i](x)
+        step += 1
+        if step in checkpoints:
+            cv_n = compute_cv(F.normalize(x, dim=-1), N_CV)
+            ed = eff_dim(x)
+            cos = (F.normalize(pts, dim=-1) * F.normalize(x, dim=-1)).sum(-1).mean().item()
+            print(f"  {step:>6} {'relay':>8} {cv_n:>8.4f} {ed:>8.1f} {cos:>10.6f}")
+
+
+# ── TEST 4: Throughput comparison ──
+print(f"\n{'━'*90}")
+print("TEST 4: Throughput")
+print(f"{'━'*90}")
+
+print(f"  {'arch':>12} {'ms':>8} {'params':>10}")
+
+for name, builder in configs.items():
+    module = builder()
+    np_ = sum(p.numel() for p in module.parameters())
+
+    # Warmup
+    for _ in range(10):
+        with torch.no_grad():
+            if name == 'relay':
+                _ = module(pts)
+            else:
+                _ = module(pts.unsqueeze(0))
+    torch.cuda.synchronize()
+
+    t0 = time.time()
+    for _ in range(200):
+        with torch.no_grad():
+            if name == 'relay':
+                _ = module(pts)
+            else:
+                _ = module(pts.unsqueeze(0))
+    torch.cuda.synchronize()
+    ms = (time.time() - t0) / 200 * 1000
+    print(f"  {name:>12} {ms:>8.2f} {np_:>10,}")
+
+
+# ── TEST 5: Clustered input — all architectures ──
+print(f"\n{'━'*90}")
+print("TEST 5: Clustered input (10 clusters, d=128)")
+print(f"{'━'*90}")
+
+centroids = F.normalize(torch.randn(10, D), dim=-1).to(DEVICE)
+assignments = torch.randint(0, 10, (N,))
+
+print(f"  {'spread':>8} {'CV_base':>8} {'vanilla':>8} {'rope_std':>8} "
+      f"{'rope_ntk':>8} {'relay':>8}")
+
+for spread in [0.1, 0.3, 0.5, 1.0]:
+    pts_c = F.normalize(centroids[assignments] +
+                        torch.randn(N, D, device=DEVICE) * spread, dim=-1)
+    cv_b = compute_cv(pts_c, N_CV)
+
+    row = f"  {spread:>8.1f} {cv_b:>8.4f}"
+    for name, builder in configs.items():
+        module = builder()
+        with torch.no_grad():
+            if name == 'relay':
+                out = module(pts_c)
+            else:
+                out = module(pts_c.unsqueeze(0)).squeeze(0)
+        cv = compute_cv(F.normalize(out, dim=-1), N_CV)
+        row += f" {cv:>8.4f}"
+    print(row)
+
+
+# ── TEST 6: RoPE frequency analysis ──
+print(f"\n{'━'*90}")
+print("TEST 6: RoPE base frequency sweep")
+print(f"  Does the rotation frequency affect geometric preservation?")
+print(f"{'━'*90}")
+
+print(f"  {'base':>10} {'CV_n':>8} {'cos_orig':>10} {'eff_dim':>8}")
+
+for base in [100, 500, 1000, 5000, 10000, 50000, 100000, 500000]:
+    module = RoPEAttnBlock(D, 8, 'standard', base).to(DEVICE)
+    with torch.no_grad():
+        out = module(pts.unsqueeze(0)).squeeze(0)
+    cv_n = compute_cv(F.normalize(out, dim=-1), N_CV)
+    cos = (F.normalize(pts, dim=-1) * F.normalize(out, dim=-1)).sum(-1).mean().item()
+    ed = eff_dim(out)
+    print(f"  {base:>10} {cv_n:>8.4f} {cos:>10.6f} {ed:>8.1f}")
+
+
+# ── TEST 7: NTK scale factor sweep ──
+print(f"\n{'━'*90}")
+print("TEST 7: NTK scale factor sweep (base=10000)")
+print(f"{'━'*90}")
+
+print(f"  {'scale':>8} {'CV_n':>8} {'cos_orig':>10} {'eff_dim':>8}")
+
+for scale in [1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0]:
+    module = RoPEAttnBlock(D, 8, 'ntk', 10000, scale).to(DEVICE)
+    with torch.no_grad():
+        out = module(pts.unsqueeze(0)).squeeze(0)
+    cv_n = compute_cv(F.normalize(out, dim=-1), N_CV)
+    cos = (F.normalize(pts, dim=-1) * F.normalize(out, dim=-1)).sum(-1).mean().item()
+    ed = eff_dim(out)
+    print(f"  {scale:>8.1f} {cv_n:>8.4f} {cos:>10.6f} {ed:>8.1f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# SUMMARY
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'='*90}")
+print("SUMMARY — cos_to_orig at depth 16")
+print(f"{'='*90}")
+print(f"""
+  Compare the depth-16 cos_to_orig from Test 2 across all architectures:
+  
+    vanilla attention:     (see Test 2)
+    RoPE standard:         (see Test 2)
+    RoPE NTK:              (see Test 2)
+    constellation relay:   (see Test 2)
+    
+  And the interleaved results from Test 3:
+    vanilla + relay:       (see Test 3)
+    rope_std + relay:      (see Test 3)
+    rope_ntk + relay:      (see Test 3)
+""")
+print(f"{'='*90}")
+print("DONE")
+print(f"{'='*90}")