training note + addr_proj: the converging configuration (rotary+addr_proj+train_codebooks), measured; frozen-all collapses in training

splat_attention.py

@@ -61,12 +61,21 @@
 #   y = attn(x, key_padding_mask=kpm)   # kpm: (B, L) True = pad
 #   python splat_attention.py    # runs the demo + a small speed bench
 #
+# TRAINING NOTE (measured 2026-08-06/07, 500k-caption encoder screens):
+#   the all-frozen configuration COLLAPSES when trained inside a trunk
+#   (representation erank ~5) — parameter-free routing deforms token
+#   states into address basins. The configuration that CONVERGES:
+#     SplatAttention(..., rotary=True, addr_proj=True,
+#                    train_codebooks=True)
+#   (routing-owned parameters: a learned address frame + living
+#   codebooks). It reaches ~90% of a standard block's training-gauge
+#   performance at matched small budget and was still climbing at
+#   cutoff — functional, slower to organize, endpoint parity unproven.
+#   Frozen-everything remains fine for INFERENCE-style play and the
+#   static properties above.
+#
 # Status: research prototype. Trained-at-scale results pending; treat
 # every number above as what it is — a measurement on the stated probe.
-#
-# License: MIT
-# Author: AbstractPhil + Claude
-# Home: https://huggingface.co/AbstractPhil/aleph-splat-0/
 # =========================================================================
 
 import math
@@ -115,7 +124,7 @@ class SplatAttention(nn.Module):
     def __init__(self, d_model, M=64, K=8, tau=0.1, dropout=0.0,
                  rotary=True, global_frac=0.0, overlap=1.5,
                  sigma_floor=0.0, head_gates=False, train_centers=False,
-                 train_codebooks=False, mchunk=16,
+                 train_codebooks=False, addr_proj=False, mchunk=16,
                  checkpoint_chunks=False):
         super().__init__()
         self.M, self.K, self.tau = M, K, tau
@@ -131,6 +140,13 @@ class SplatAttention(nn.Module):
             self.register_buffer("codebook", book)
         self.w_v = nn.Linear(d_model, d_model)
         self.w_o = nn.Linear(d_model, d_model)
+        if addr_proj:
+            # learned address frame, born at identity — routing-owned
+            # parameters (see TRAINING note in the module docstring)
+            self.w_a = nn.Linear(d_model, d_model, bias=False)
+            nn.init.eye_(self.w_a.weight)
+        else:
+            self.w_a = None
         self.drop = nn.Dropout(dropout)
         if train_centers:
             self.center_off = nn.Parameter(torch.zeros(M))
@@ -187,7 +203,8 @@ class SplatAttention(nn.Module):
 
     def forward(self, x, key_padding_mask=None):
         B, L, d = x.shape
-        xn = F.normalize(x, dim=-1)
+        xa = self.w_a(x) if self.w_a is not None else x
+        xn = F.normalize(xa, dim=-1)
         if self.rotary:
             pos = torch.arange(L, device=x.device, dtype=torch.float32)
             xn = F.normalize(rope_rotate(xn.float(), pos),
@@ -216,7 +233,6 @@ class SplatAttention(nn.Module):
         return self.drop(self.w_o(out))
 
 
-
 def _demo():
     torch.manual_seed(0)
     dev = "cuda" if torch.cuda.is_available() else "cpu"
@@ -232,7 +248,7 @@ def _demo():
     if dev == "cuda":
         import time
         mha = nn.MultiheadAttention(256, 8, batch_first=True).to(dev)
-        for L, B in [(128, 256), (256, 32), (512, 32), (768, 32), (1024, 32), (1280, 32), (2048, 32), (4096, 32), (8192, 32)]:
+        for L, B in [(128, 32), (2048, 2)]:
             xx = torch.randn(B, L, 256, device=dev, requires_grad=True)
             def t(fn, n=10):
                 for _ in range(3):
@@ -251,6 +267,5 @@ def _demo():
           "head_gates=True for learnable attenuation)")
 
 
-
 if __name__ == "__main__":
     _demo()