Upload chimera.py with huggingface_hub

chimera.py

@@ -0,0 +1,215 @@
+"""
+chimera.py -- the ChimeraBlock: the grand-finale channel mixer that fuses EVERY project
+in the family into one block. A creature made of parts of many beasts.
+
+Per token, three PHYSICS CORES each propose a candidate update of the hidden state:
+
+  * KuramotoCore  -- coupled phase oscillators (Quazimoto): mean-field Kuramoto with
+                     learnable frustration; a few Euler steps; readout [cos, sin].
+  * GrowthCore    -- Neighbour-Sensing fungal growth (Mycel): tips in a bounded latent
+                     region sense a low-rank density field and steer (negative autotropism).
+  * WaveCore      -- Wheeler-DeWitt wave (Wheeler): K minisuperspace modes under a
+                     LORENTZIAN supermetric, leapfrog wave steps; exposes the Hamiltonian
+                     constraint <H^2> so the block can be pressured onto H Psi = 0.
+
+Then the GRPO GROUP-RELATIVE SELECTION (grpo_lm) is the META-MIXER: an internal critic
+scores the three candidates, group-relative advantages A_g = (r_g - mean)/std decide the
+winner RELATIVE to its peers, the mixing weights are CLIPPED around uniform (the PPO clip)
+and ANCHORED back toward uniform by beta (the KL-to-reference anchor). The output is the
+selected convex combination -- so the network LEARNS WHICH LAW OF PHYSICS to apply to each
+token. Fractal phase (fractal.py) seeds all three cores; Elo attention lives in the
+attention block. Every idea we designed, in one place, behind one family gate.
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from family import soft_clamp, RMSNorm
+import instrument as _viz
+
+
+# --------------------------------------------------------------------------------------
+# physics cores -- each takes normed hidden h [B,T,d] (+ optional fractal seed) and returns
+# a candidate update [B,T,d]. Compact readouts (feat -> d -> d) keep the three-core cost sane.
+# --------------------------------------------------------------------------------------
+class KuramotoCore(nn.Module):
+    """Quazimoto: a single mean-field Kuramoto ring with learnable frustration alpha."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.N = cfg.chim_osc
+        self.steps, self.dt = cfg.chim_osc_steps, cfg.osc_dt
+        self.to_theta = nn.Linear(cfg.d_model, self.N)
+        self.to_omega = nn.Linear(cfg.d_model, self.N)
+        self.k_coupling = nn.Parameter(torch.tensor(1.0))     # global coupling strength
+        self.alpha = nn.Parameter(torch.zeros(1))             # Sakaguchi frustration
+        self.read = nn.Sequential(nn.Linear(2 * self.N, cfg.d_model), nn.GELU(),
+                                  nn.Linear(cfg.d_model, cfg.d_model))
+
+    def forward(self, h, seed=None):
+        theta = self.to_theta(h)
+        if seed is not None:
+            theta = theta + seed[..., :self.N]
+        omega = torch.tanh(self.to_omega(h))
+        a = self.alpha
+        for _ in range(self.steps):
+            c, s = torch.cos(theta), torch.sin(theta)
+            zc = c.mean(-1, keepdim=True)                     # mean-field order parameter
+            zs = s.mean(-1, keepdim=True)
+            # K * Im( e^{i(alpha - theta)} * z ) = K*(sin(a-θ)zc + cos(a-θ)zs)
+            coupling = self.k_coupling * (torch.sin(a - theta) * zc + torch.cos(a - theta) * zs)
+            theta = theta + self.dt * (omega + coupling)
+        return self.read(torch.cat([torch.cos(theta), torch.sin(theta)], dim=-1))
+
+
+class GrowthCore(nn.Module):
+    """Mycel: Neighbour-Sensing tips growing in a bounded latent region, sensing a
+    low-rank density field (O(N*F)) and steering away from their own density."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.N, self.pd, self.F = cfg.chim_tips, cfg.chim_pos_dim, cfg.chim_centers
+        self.steps, self.dt, self.bound = cfg.chim_growth_steps, cfg.chim_growth_dt, cfg.osc_bound
+        self.to_pos = nn.Linear(cfg.d_model, self.N * self.pd)
+        self.to_dir = nn.Linear(cfg.d_model, self.N * self.pd)
+        self.centers = nn.Parameter(torch.randn(self.F, self.pd) * 0.5)
+        self.log_persist = nn.Parameter(torch.tensor(1.4))
+        self.tropism = nn.Parameter(torch.zeros(1))
+        self.log_bw = nn.Parameter(torch.zeros(1))
+        self.read = nn.Sequential(nn.Linear(self.N * (2 * self.pd + 1), cfg.d_model), nn.GELU(),
+                                  nn.Linear(cfg.d_model, cfg.d_model))
+
+    def _clamp(self, p):
+        return self.bound * torch.tanh(p / self.bound)
+
+    def _cdist2(self, p, c):
+        p2 = (p * p).sum(-1, keepdim=True)
+        c2 = (c * c).sum(-1)
+        return (p2 + c2 - 2.0 * (p @ c.t())).clamp(min=0.0)
+
+    def forward(self, h, seed=None):
+        B, T, _ = h.shape
+        p = self._clamp(self.to_pos(h).view(B, T, self.N, self.pd))
+        if seed is not None:
+            p = self._clamp(p + seed[..., :self.N * self.pd].view(B, T, self.N, self.pd))
+        v = self.to_dir(h).view(B, T, self.N, self.pd)
+        pers, trop = torch.sigmoid(self.log_persist), torch.tanh(self.tropism)
+        bw = F.softplus(self.log_bw).clamp(min=1e-3)
+        for _ in range(self.steps):
+            K = torch.exp(-self._cdist2(p, self.centers) / bw)
+            w = K.mean(-2, keepdim=True) * K
+            away = p * w.sum(-1, keepdim=True) - torch.matmul(w, self.centers)
+            v = pers * v + trop * away
+            p = self._clamp(p + self.dt * v)
+        dens = torch.exp(-self._cdist2(p, self.centers) / bw).mean(-1, keepdim=True)
+        feat = torch.cat([p, v, dens], dim=-1).flatten(2)
+        return self.read(feat)
+
+
+class WaveCore(nn.Module):
+    """Wheeler: K minisuperspace modes under a learnable LORENTZIAN DeWitt supermetric,
+    leapfrog wave dynamics, curvature potential. Exposes the Hamiltonian constraint."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.K = cfg.wdw_modes
+        self.steps, self.dt = cfg.wdw_steps, cfg.wdw_dt
+        self.to_psi = nn.Linear(cfg.d_model, self.K)
+        self.to_pi = nn.Linear(cfg.d_model, self.K)
+        self.to_curv = nn.Linear(cfg.d_model, self.K)
+        self.ginv_raw = nn.Parameter(torch.zeros(self.K, self.K))
+        sig = torch.ones(self.K); sig[0] = -1.0
+        self.register_buffer("signature", sig)
+        self.log_lapse = nn.Parameter(torch.zeros(1))
+        self.read = nn.Sequential(nn.Linear(2 * self.K + 2, cfg.d_model), nn.GELU(),
+                                  nn.Linear(cfg.d_model, cfg.d_model))
+        self.last_H = None
+
+    def _supermetric(self):
+        A = 0.5 * (self.ginv_raw + self.ginv_raw.t())
+        return A + torch.diag(self.signature)
+
+    def forward(self, h, seed=None):
+        psi = self.to_psi(h)
+        if seed is not None:
+            psi = psi + seed[..., :self.K]
+        pi = self.to_pi(h)
+        r = self.to_curv(h)
+        Ginv = self._supermetric()
+        dt = self.dt * F.softplus(self.log_lapse).clamp(max=4.0)
+        pi = pi - 0.5 * dt * (r * psi)
+        for j in range(self.steps):
+            psi = psi + dt * torch.matmul(pi, Ginv.t())
+            pi = pi - (0.5 if j == self.steps - 1 else 1.0) * dt * (r * psi)
+        H = 0.5 * (torch.matmul(pi, Ginv.t()) * pi).sum(-1) + 0.5 * (r * psi * psi).sum(-1)
+        self.last_H = (H ** 2).mean() if self.training else None
+        t_comp, space_norm = psi[..., :1], psi[..., 1:].norm(dim=-1, keepdim=True)
+        return self.read(torch.cat([psi, pi, t_comp, space_norm], dim=-1))
+
+
+# --------------------------------------------------------------------------------------
+# the Chimera block: council of the three cores, combined by GRPO group-relative selection
+# --------------------------------------------------------------------------------------
+class ChimeraBlock(nn.Module):
+    """Drop-in for QuazimotoBlock.forward(x, ring_ctl, phase_seed). Runs three physics
+    cores and GRPO-selects among them (critic -> group-relative advantage -> clip -> anchor),
+    behind one family gate. Exposes last_constraint (the wave core's <H^2>)."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.norm = RMSNorm(cfg.d_model)
+        self.cores = nn.ModuleList([KuramotoCore(cfg), GrowthCore(cfg), WaveCore(cfg)])
+        self.G = len(self.cores)
+        self.critic = nn.Linear(cfg.d_model, 1)               # scores each candidate (the verifier)
+        self.adv_clip, self.temp = cfg.chim_adv_clip, cfg.chim_select_temp
+        self.clip, self.anchor = cfg.chim_clip, cfg.chim_anchor
+        self.drop = nn.Dropout(cfg.dropout)
+        go = math.atanh(min(cfg.gate_init_open, 0.9)) if cfg.gate_init_open > 0 else 0.0
+        self.gate = nn.Parameter(torch.full((1,), go))
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+        if cfg.use_fractal_phase_seed:
+            self.seed_gate = nn.Parameter(torch.zeros(1))     # zero-init -> no fractal seed at start
+        self.last_constraint = None
+        self.last_select = None                               # mean selection weights (for viz)
+
+    def forward(self, x, ring_ctl=None, phase_seed=None):
+        cfg = self.cfg
+        h = self.norm(x)
+        seed = None
+        if phase_seed is not None and cfg.use_fractal_phase_seed:
+            seed = torch.tanh(self.seed_gate) * phase_seed
+
+        cands = torch.stack([core(h, seed) for core in self.cores], dim=2)   # [B,T,G,d]
+
+        # ---- GRPO group-relative selection over the G candidates ----
+        r = self.critic(cands).squeeze(-1)                    # [B,T,G] critic reward per core
+        mean = r.mean(-1, keepdim=True)
+        std = r.std(-1, keepdim=True)
+        adv = ((r - mean) / (std + 1e-4)).clamp(-self.adv_clip, self.adv_clip)   # group-relative
+        sel = torch.softmax(adv / max(self.temp, 1e-6), dim=-1)                  # winners get mass
+        u = 1.0 / self.G                                                          # uniform = "old policy"
+        sel = sel.clamp(u * (1.0 - self.clip), u * (1.0 + self.clip))            # PPO clip around uniform
+        sel = sel / sel.sum(-1, keepdim=True)
+        w = (1.0 - self.anchor) * sel + self.anchor * u                         # KL-to-uniform anchor
+        out = (w.unsqueeze(-1) * cands).sum(2)                                   # [B,T,d]
+        out = self.drop(soft_clamp(out * torch.tanh(self.gate), cfg.osc_bound))
+
+        # wave core's Hamiltonian constraint -> trunk aux loss (H Psi = 0 pressure)
+        wave = self.cores[2]
+        self.last_constraint = wave.last_H
+        self.last_select = w.detach().mean((0, 1)) if not self.training else None
+
+        rec = _viz.get_rec()
+        if rec is not None and rec.enabled:                   # live-viz: which law won this token
+            wsel = w[0, -1].tolist()                          # [G] selection weights
+            hval = float(wave.last_H) if wave.last_H is not None else 0.0
+            rec.log_ring([hval], wsel, wsel)
+            rec.log_quaz_norm(out[0, -1].norm().item())
+            rec.flush_spec()
+        return out