AbstractPhil/mini-beatrix-1
12.4k
1"""The anchor governor — min-separation projection (ROUND 5f, 2026-08-25).2 3Measured law (claude-mind, canon/splat_aleph_battery.md ROUND 5e/5f): anchors4crowded past the separation floor flutter under the positional sweep (the5static logit gap is quadratic in separation, the oscillation linear), and the6crowding tail tracks seed fragility. The governor holds the space open FROM7BIRTH — the projection is preventive, not curative (exact lemma: separating an8already-crowded book leaves its flip pattern bit-identical for content already9associated with it).10 11Form: a post-optimizer-step PROJECTION, never an auxiliary loss — this repo's12standing law ("no balance machinery; pressure stays out of the task gradient",13model/bank.py) is respected because the governor never touches the gradient.14Hub and head codebooks live natively in their own D-dim address space, so the15plain projective form applies (|cos| on RP^(D-1): the read is exactly16invariant under a -> -a). Identity when slack: the batched check is one17matmul per codebook and nothing is written unless a pair violates.18 19Verified properties (2026-08-25, gov25m rung): fires only when needed20(bit-identical runs when slack), census tail cut 5.4x, fragile-seed function21+.007, stable-seed +.002 — zero cost, no parameters.22"""23from __future__ import annotations24 25import math26 27import torch28 29from .address import AlephAddress30 31__all__ = ["minsep_project_", "govern_model"]32 33 34@torch.no_grad()35def minsep_project_(codebook: torch.nn.Parameter, theta_min_deg: float,36 max_iters: int = 8) -> int:37 """Projective min-separation on one (K, D) codebook, in place.38 39 Bound: |cos(a_i, a_j)| <= cos(theta_min) for every pair. Exact symmetric40 geodesic push-apart, largest violation first, projected to theta_min +41 0.5 deg (the exact-bound push converges asymptotically under fp32 and42 can land ~1e-4 over). Returns the number of pair-pushes performed —43 0 means the call was a pure read (identity, zero writes).44 """45 c_max = math.cos(math.radians(theta_min_deg))46 A = torch.nn.functional.normalize(codebook.data.float(), dim=-1)47 hits = 048 changed = False49 for _ in range(max_iters):50 G = A @ A.T51 G.fill_diagonal_(0)52 viol = (G.abs() > c_max).nonzero()53 viol = viol[viol[:, 0] < viol[:, 1]]54 if not len(viol):55 break56 order = G.abs()[viol[:, 0], viol[:, 1]].argsort(descending=True)57 for i, j in viol[order].tolist():58 g = float(A[i] @ A[j])59 if abs(g) <= c_max:60 continue61 s = 1.0 if g >= 0 else -1.062 b = s * A[j]63 gam = math.acos(max(-1.0, min(1.0, float(A[i] @ b))))64 if gam < 1e-6: # exact duplicate: nudge first65 b = torch.nn.functional.normalize(66 b + 1e-3 * torch.randn_like(b), dim=-1)67 gam = math.acos(max(-1.0, min(1.0, float(A[i] @ b))))68 delta = (math.radians(theta_min_deg + 0.5) - gam) / 269 ai = A[i].clone()70 ti = (b - math.cos(gam) * ai) / math.sin(gam)71 tj = (ai - math.cos(gam) * b) / math.sin(gam)72 A[i] = torch.nn.functional.normalize(73 math.cos(delta) * ai - math.sin(delta) * ti, dim=-1)74 A[j] = s * torch.nn.functional.normalize(75 math.cos(delta) * b - math.sin(delta) * tj, dim=-1)76 hits += 177 changed = True78 if changed:79 codebook.data.copy_(A.to(codebook.dtype))80 return hits81 82 83def raw_block(wrap):84 """The model's own Block under any depth of adapter wrappers (amoe85 BlockWithAdapter nests one wrapper per attached arm: .block.block...)."""86 while hasattr(wrap, "block"):87 wrap = wrap.block88 return wrap89 90 91def _hub_addresses(attn):92 """Every AlephAddress a hub carries (single- or multi-constellation)."""93 if hasattr(attn, "consts"):94 return [c.addr for c in attn.consts]95 addr = getattr(attn, "addr", None)96 return [addr] if isinstance(addr, AlephAddress) else []97 98 99@torch.no_grad()100def govern_model(model, theta_min_deg: float,101 include: tuple = ("hub", "head")) -> int:102 """Batched slack check + projection across a model's governed codebooks.103 104 Slack path costs one small matmul per codebook and writes nothing.105 Wrapper-aware (amoe BlockWithAdapter moves the block to .block).106 `include`: "hub" = attention codebooks, "head" = DualHead, "bank" =107 dispatch anchors (K=experts in d_model — never crowds in practice, off108 by default).109 """110 c_max = math.cos(math.radians(theta_min_deg))111 books: list[torch.nn.Parameter] = []112 if "hub" in include:113 for wrap in getattr(model, "blocks", []):114 blk = raw_block(wrap)115 if getattr(blk, "is_hub", False):116 books += [a.codebook for a in _hub_addresses(blk.attn)]117 if "head" in include and hasattr(model, "head"):118 addr = getattr(model.head, "addr", None)119 if isinstance(addr, AlephAddress):120 books.append(addr.codebook)121 if "bank" in include:122 for wrap in getattr(model, "blocks", []):123 blk = raw_block(wrap)124 addr = getattr(getattr(blk, "bank", None), "addr", None)125 if isinstance(addr, AlephAddress):126 books.append(addr.codebook)127 # Batched slack check (2026-08-26): group same-shape books into ONE bmm128 # and ONE device sync — the per-book loop was 500+ tiny matmuls each129 # with a .max() sync on the v2 craft. Only violating books ever enter130 # the Python projection.131 hits = 0132 by_shape: dict = {}133 for cb in books:134 by_shape.setdefault(tuple(cb.shape), []).append(cb)135 for shape, group in by_shape.items():136 A = torch.nn.functional.normalize(137 torch.stack([cb.data for cb in group]).float(), dim=-1)138 G = torch.bmm(A, A.transpose(1, 2)).abs() # (n_books, K, K)139 K = shape[0]140 eye = torch.eye(K, device=G.device, dtype=torch.bool)141 G = G.masked_fill(eye, 0)142 worst = G.amax(dim=(1, 2)) # one sync for all143 for idx in (worst > c_max).nonzero().flatten().tolist():144 hits += minsep_project_(group[idx], theta_min_deg)145 return hits146 