| |
| """ |
| Constellation Relay v2 β Vectorized |
| ===================================== |
| All patches triangulated in a single batched matmul. |
| Gates init near zero (relay off at init, learns to open). |
| |
| Architecture: |
| input (B, D) β LN β reshape (B, P, d) |
| β L2 normalize each patch to S^(d-1) |
| β batched tri: (B, P, d) Γ (P, d, A*phases) β (B, P, A*phases) |
| β patchwork MLP per patch (shared or independent) |
| β gated residual (gate init β 0.05) |
| β reshape (B, D) + skip connection |
| """ |
|
|
| 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) |
|
|
|
|
| def uniform_hypersphere_init(n, d): |
| if n <= d: |
| M = torch.randn(d, n) |
| Q, _ = torch.linalg.qr(M) |
| return Q.T.contiguous() |
| else: |
| M = torch.randn(d, d) |
| Q, _ = torch.linalg.qr(M) |
| basis = Q.T |
| extra = F.normalize(torch.randn(n - d, d), dim=-1) |
| vecs = torch.cat([basis, extra], dim=0) |
| for _ in range(200): |
| sim = vecs @ vecs.T |
| sim.fill_diagonal_(-2.0) |
| nn_idx = sim.argmax(dim=1) |
| vecs = F.normalize(vecs - 0.05 * vecs[nn_idx], dim=-1) |
| return vecs |
|
|
|
|
| |
| |
| |
|
|
| class ConstellationRelay(nn.Module): |
| """ |
| Fully vectorized constellation relay. |
| |
| All P patches triangulated in one einsum. |
| Gates init near zero β relay is transparent at init. |
| """ |
| 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): |
| """Per-anchor drift: (P, A) in radians.""" |
| 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): |
| """Slerp all anchors: (P, A, d).""" |
| 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 _fp8_or_native(self, patches_n, anchors_n): |
| """ |
| Batched triangulation: patches (B, P, d) Γ anchors (P, A, d) β (B, P, A) |
| Uses fp8 if available, otherwise native. |
| """ |
| if HAS_FP8: |
| |
| amax_p = patches_n.abs().amax().clamp(min=1e-12) |
| amax_a = anchors_n.abs().amax().clamp(min=1e-12) |
| sp = torch.finfo(torch.float8_e4m3fn).max / amax_p |
| sa = torch.finfo(torch.float8_e4m3fn).max / amax_a |
| p8 = (patches_n * sp).to(torch.float8_e4m3fn).to(patches_n.dtype) / sp |
| a8 = (anchors_n * sa).to(torch.float8_e4m3fn).to(anchors_n.dtype) / sa |
| |
| return torch.einsum('bpd,pad->bpa', p8, a8) |
| else: |
| return torch.einsum('bpd,pad->bpa', patches_n, anchors_n) |
|
|
| def triangulate_multiphase(self, patches_n): |
| """ |
| Multi-phase stroboscope, fully vectorized. |
| patches_n: (B, P, d) normalized |
| Returns: (B, P, A * n_phases) |
| """ |
| 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 = self._fp8_or_native(patches_n, anchors_t) |
| tris.append(1.0 - cos) |
| return torch.cat(tris, dim=-1) |
|
|
| def forward(self, x, return_diagnostics=False): |
| 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) |
|
|
| |
| tri = self.triangulate_multiphase(patches_n) |
|
|
| |
| |
| 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) |
| result = x + out |
|
|
| if return_diagnostics: |
| |
| anchors_n = F.normalize(self.anchors, dim=-1) |
| cos_all = torch.einsum('bpd,pad->bpa', patches_n, anchors_n) |
| cos64 = cos_all.to(torch.float64) |
| top2 = cos64.topk(2, dim=-1) |
| nearest_cos = top2.values[:, :, 0].float().mean().item() |
| margin = (top2.values[:, :, 0] - top2.values[:, :, 1]).float().mean().item() |
| drift = self.drift() |
| diag = { |
| 'nearest_cos': nearest_cos, |
| 'margin': margin, |
| 'drift_mean': drift.mean().item(), |
| 'drift_max': drift.max().item(), |
| 'gate_mean': self.gates.sigmoid().mean().item(), |
| 'gate_min': self.gates.sigmoid().min().item(), |
| 'gate_max': self.gates.sigmoid().max().item(), |
| } |
| return result, diag |
| return result |
|
|
| def get_diagnostics(self): |
| drift = self.drift().detach().cpu() |
| gates = self.gates.sigmoid().detach().cpu() |
| return { |
| 'drift_mean': drift.mean().item(), |
| 'drift_std': drift.std().item(), |
| 'gate_mean': gates.mean().item(), |
| 'gate_std': gates.std().item(), |
| } |
|
|
|
|
| |
| |
| |
|
|
| class SelfAttnBlock(nn.Module): |
| 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) |
|
|
|
|
| |
| |
| |
|
|
| N = 2000 |
| D = 128 |
| N_CV = 2000 |
|
|
| print("=" * 80) |
| print("CONSTELLATION RELAY v2 β VECTORIZED + COLD GATES") |
| print(f" Input dim: {D} = {D//16} patches Γ 16") |
| print(f" Device: {DEVICE}") |
| print(f" fp8: {'hardware' if HAS_FP8 else 'simulated'}") |
| print("=" * 80) |
|
|
|
|
| |
| print(f"\n{'β'*80}") |
| print("TEST 1: Single relay pass β CV preservation + gate state") |
| print(f"{'β'*80}") |
|
|
| relay = ConstellationRelay(D, patch_dim=16, n_anchors=16, n_phases=3, pw_hidden=32).to(DEVICE) |
| n_params_relay = sum(p.numel() for p in relay.parameters()) |
| print(f" Relay params: {n_params_relay:,}") |
| print(f" Gate init: sigmoid({relay.gates[0].item():.1f}) = {relay.gates[0].sigmoid().item():.4f}") |
|
|
| pts = uniform_sphere(N, D).to(DEVICE) |
| cv_before = compute_cv(pts, n_samples=N_CV) |
| ed_before = eff_dim(pts) |
|
|
| with torch.no_grad(): |
| out, diag = relay(pts, return_diagnostics=True) |
|
|
| out_n = F.normalize(out, dim=-1) |
| cv_after = compute_cv(out, n_samples=N_CV) |
| cv_normed = compute_cv(out_n, n_samples=N_CV) |
| ed_after = eff_dim(out) |
| cos_orig = (F.normalize(pts, dim=-1) * F.normalize(out, dim=-1)).sum(-1).mean().item() |
|
|
| print(f" Before: CV={cv_before:.4f} eff_dim={ed_before:.1f}") |
| print(f" After: CV={cv_after:.4f} eff_dim={ed_after:.1f}") |
| print(f" Normed: CV={cv_normed:.4f}") |
| print(f" cos_to_orig: {cos_orig:.6f}") |
| print(f" Diagnostics: cos={diag['nearest_cos']:.4f} margin={diag['margin']:.4f} " |
| f"drift={diag['drift_mean']:.4f} gate={diag['gate_mean']:.4f}") |
|
|
|
|
| |
| print(f"\n{'β'*80}") |
| print("TEST 2: Depth stability β relay vs attention vs interleaved") |
| print(f"{'β'*80}") |
|
|
| pts = uniform_sphere(N, D).to(DEVICE) |
|
|
| |
| print(f"\n A) Relay-only:") |
| print(f" {'cycle':>6} {'CV':>8} {'CV_n':>8} {'eff_d':>8} {'cos_orig':>10} {'gate':>6}") |
|
|
| relay_stack = nn.ModuleList([ |
| ConstellationRelay(D, 16, 16, 3, 32).to(DEVICE) for _ in range(16)]) |
|
|
| x = pts.clone() |
| for i, r in enumerate(relay_stack): |
| with torch.no_grad(): |
| x = r(x) |
| if (i + 1) in [1, 2, 4, 8, 12, 16]: |
| 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() |
| g = r.gates.sigmoid().mean().item() |
| print(f" {i+1:>6} {cv:>8.4f} {cv_n:>8.4f} {ed:>8.1f} {cos:>10.6f} {g:>6.3f}") |
|
|
| |
| print(f"\n B) Attention-only:") |
| print(f" {'cycle':>6} {'CV':>8} {'CV_n':>8} {'eff_d':>8} {'cos_orig':>10}") |
|
|
| attn_stack = nn.ModuleList([ |
| SelfAttnBlock(D, n_heads=8).to(DEVICE) for _ in range(16)]) |
|
|
| x = pts.clone() |
| for i, layer in enumerate(attn_stack): |
| with torch.no_grad(): |
| x = layer(x.unsqueeze(0)).squeeze(0) |
| if (i + 1) in [1, 2, 4, 8, 12, 16]: |
| 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}") |
|
|
| |
| print(f"\n C) Interleaved (attn β relay β ...):") |
| print(f" {'step':>6} {'type':>6} {'CV_n':>8} {'eff_d':>8} {'cos_orig':>10}") |
|
|
| attn_il = nn.ModuleList([SelfAttnBlock(D, 8).to(DEVICE) for _ in range(8)]) |
| relay_il = nn.ModuleList([ConstellationRelay(D, 16, 16, 3, 32).to(DEVICE) for _ in range(8)]) |
|
|
| x = pts.clone() |
| step = 0 |
| for i in range(8): |
| with torch.no_grad(): |
| x = attn_il[i](x.unsqueeze(0)).squeeze(0) |
| step += 1 |
| if step in [1, 2, 4, 8, 12, 16]: |
| 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':>6} {cv_n:>8.4f} {ed:>8.1f} {cos:>10.6f}") |
|
|
| with torch.no_grad(): |
| x = relay_il[i](x) |
| step += 1 |
| if step in [1, 2, 4, 8, 12, 16]: |
| 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':>6} {cv_n:>8.4f} {ed:>8.1f} {cos:>10.6f}") |
|
|
|
|
| |
| print(f"\n{'β'*80}") |
| print("TEST 3: Clustered input β relay vs attention") |
| print(f"{'β'*80}") |
|
|
| centroids = F.normalize(torch.randn(10, D), dim=-1).to(DEVICE) |
| assignments = torch.randint(0, 10, (N,)) |
|
|
| print(f" {'spread':>8} {'CV_before':>10} {'CV_relay':>10} {'CV_attn':>10} " |
| f"{'Ξ_relay':>8} {'Ξ_attn':>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) |
|
|
| r = ConstellationRelay(D, 16, 16, 3, 32).to(DEVICE) |
| with torch.no_grad(): |
| out_r = r(pts_c) |
| cv_r = compute_cv(F.normalize(out_r, dim=-1), N_CV) |
|
|
| a = SelfAttnBlock(D, 8).to(DEVICE) |
| with torch.no_grad(): |
| out_a = a(pts_c.unsqueeze(0)).squeeze(0) |
| cv_a = compute_cv(F.normalize(out_a, dim=-1), N_CV) |
|
|
| print(f" {spread:>8.1f} {cv_b:>10.4f} {cv_r:>10.4f} {cv_a:>10.4f} " |
| f"{cv_r-cv_b:>+8.4f} {cv_a-cv_b:>+8.4f}") |
|
|
|
|
| |
| print(f"\n{'β'*80}") |
| print("TEST 4: Throughput β vectorized relay vs attention") |
| print(f"{'β'*80}") |
|
|
| relay_b = ConstellationRelay(D, 16, 16, 3, 32).to(DEVICE) |
| attn_b = SelfAttnBlock(D, 8).to(DEVICE) |
| pts_b = uniform_sphere(N, D).to(DEVICE) |
|
|
| |
| for _ in range(10): |
| with torch.no_grad(): |
| _ = relay_b(pts_b) |
| _ = attn_b(pts_b.unsqueeze(0)) |
| torch.cuda.synchronize() |
|
|
| t0 = time.time() |
| for _ in range(200): |
| with torch.no_grad(): |
| _ = relay_b(pts_b) |
| torch.cuda.synchronize() |
| relay_ms = (time.time() - t0) / 200 * 1000 |
|
|
| t0 = time.time() |
| for _ in range(200): |
| with torch.no_grad(): |
| _ = attn_b(pts_b.unsqueeze(0)) |
| torch.cuda.synchronize() |
| attn_ms = (time.time() - t0) / 200 * 1000 |
|
|
| n_params_attn = sum(p.numel() for p in attn_b.parameters()) |
| print(f" Relay: {relay_ms:.2f}ms ({n_params_relay:,} params)") |
| print(f" Attention: {attn_ms:.2f}ms ({n_params_attn:,} params)") |
| print(f" Speedup: {attn_ms/relay_ms:.2f}Γ {'β relay faster' if relay_ms < attn_ms else 'β attn faster'}") |
|
|
|
|
| |
| print(f"\n{'β'*80}") |
| print("TEST 5: Patch count sweep") |
| print(f"{'β'*80}") |
|
|
| pts = uniform_sphere(N, D).to(DEVICE) |
| print(f" {'patches':>8} {'d_patch':>8} {'params':>10} {'CV_out':>8} {'cos_orig':>10} {'ms':>8}") |
|
|
| for n_patches in [1, 2, 4, 8, 16]: |
| pd = D // n_patches |
| if D % n_patches != 0: continue |
| r = ConstellationRelay(D, pd, min(pd, 16), 3, 32).to(DEVICE) |
| np_r = sum(p.numel() for p in r.parameters()) |
|
|
| with torch.no_grad(): _ = r(pts) |
| torch.cuda.synchronize() |
|
|
| t0 = time.time() |
| for _ in range(100): |
| with torch.no_grad(): |
| out = r(pts) |
| torch.cuda.synchronize() |
| ms = (time.time() - t0) / 100 * 1000 |
|
|
| cv = compute_cv(F.normalize(out, dim=-1), N_CV) |
| cos = (F.normalize(pts, dim=-1) * F.normalize(out, dim=-1)).sum(-1).mean().item() |
| print(f" {n_patches:>8} {pd:>8} {np_r:>10,} {cv:>8.4f} {cos:>10.6f} {ms:>8.2f}") |
|
|
|
|
| |
| print(f"\n{'β'*80}") |
| print("TEST 6: Scaling β relay at d=256, 512, 768, 1024") |
| print(f"{'β'*80}") |
|
|
| print(f" {'D':>6} {'patches':>8} {'params':>10} {'CV_in':>8} {'CV_out':>8} {'cos':>10} {'ms':>8}") |
|
|
| for D_test in [256, 512, 768, 1024]: |
| pts_t = uniform_sphere(N, D_test).to(DEVICE) |
| r = ConstellationRelay(D_test, 16, 16, 3, 32).to(DEVICE) |
| np_r = sum(p.numel() for p in r.parameters()) |
| cv_in = compute_cv(pts_t, N_CV) |
|
|
| with torch.no_grad(): _ = r(pts_t) |
| torch.cuda.synchronize() |
| t0 = time.time() |
| for _ in range(50): |
| with torch.no_grad(): |
| out = r(pts_t) |
| torch.cuda.synchronize() |
| ms = (time.time() - t0) / 50 * 1000 |
|
|
| cv_out = compute_cv(F.normalize(out, dim=-1), N_CV) |
| cos = (F.normalize(pts_t, dim=-1) * F.normalize(out, dim=-1)).sum(-1).mean().item() |
| print(f" {D_test:>6} {D_test//16:>8} {np_r:>10,} {cv_in:>8.4f} {cv_out:>8.4f} " |
| f"{cos:>10.6f} {ms:>8.2f}") |
|
|
|
|
| |
| |
| |
|
|
| print(f"\n{'='*80}") |
| print("SUMMARY") |
| print(f"{'='*80}") |
| print(f""" |
| v2 changes: |
| - Vectorized: all patches in single einsum, no Python loop |
| - Gates init at sigmoid(-3.0) β 0.047 (was 0.622) |
| - Patchwork via batched einsum (P, tri, hidden) Γ (P, hidden, d) |
| |
| Architecture: {D}d β {D//16}p Γ 16d, {16}A/patch, 3-phase |
| Params: {n_params_relay:,} |
| |
| Key comparison at depth 16: |
| Relay-only cos_to_orig: (see Test 2A, cycle 16) |
| Attention-only cos_to_orig: (see Test 2B, cycle 16) |
| Interleaved cos_to_orig: (see Test 2C, step 16) |
| """) |
| print(f"{'='*80}") |
| print("DONE") |
| print(f"{'='*80}") |
