mini-beatrix-2s automodel: head.py (mission final 16.101B, alephllm 0.8.6)

head.py

@@ -0,0 +1,51 @@
+"""DualHead — the standard readout plus the L-012 aleph read, born null.
+
+    logits = W_h h + W_s s(h),        W_s == 0 at init
+
+s(h) is the signed address of a learned low-D projection of h through a
+K_h-anchor head codebook. New structure enters at zero and earns its way
+in by gradient — via WEIGHT-zero init, never gate-zero: mini-beatrix-1
+measured (8.4B tokens) that a gamma-gated branch DEADLOCKS (W_s grad is
+scaled by gamma=0, gamma grad through a random frozen readout is
+zero-mean noise), while the bank's E_out weight-zero self-started into
++2 bpb of function. gamma remains as a frozen scalar for checkpoint
+compatibility and as the ablation knob; the election gauge is ||W_s||
+and the head_aleph_off toggle.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from .address import AlephAddress
+
+
+class DualHead(nn.Module):
+    def __init__(self, d: int, vocab: int, K: int = 512, D: int = 32,
+                 tau: float = 0.1, tied_weight: nn.Parameter | None = None):
+        super().__init__()
+        self.tied = tied_weight is not None
+        if self.tied:
+            # tuple wrapper: reference the embedding table WITHOUT registering
+            # it as a duplicate parameter (state_dict must stay alias-free
+            # for safetensors export)
+            self._tied_ref = (tied_weight,)
+            self.bias = nn.Parameter(torch.zeros(vocab))
+        else:
+            self.w_h = nn.Linear(d, vocab, bias=True)
+        self.proj = nn.Linear(d, D, bias=False)
+        nn.init.orthogonal_(self.proj.weight)
+        self.addr = AlephAddress(K, D, tau)
+        self.w_s = nn.Linear(K, vocab, bias=False)
+        nn.init.zeros_(self.w_s.weight)          # weight-zero: self-starting
+        self.gamma = nn.Parameter(torch.ones(1))  # frozen; ablation knob only
+        self.gamma.requires_grad_(False)
+
+    def forward(self, h, disable_aleph: bool = False):
+        if self.tied:
+            base = h @ self._tied_ref[0].T + self.bias
+        else:
+            base = self.w_h(h)
+        if disable_aleph:
+            return base
+        return base + self.gamma * self.w_s(self.addr.signed(self.proj(h)))