CoolFace
Modelpublic

AbstractPhil/mini-beatrix-1

sourceHugging Facemitupdated 2d agoView on Hugging Face
1likes2.4kdownloads
governor.py146 linesDownload Raw Back to root
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