Spaces:
Sleeping
Sleeping
| """ | |
| wheeler.py -- the Wheeler-DeWitt channel-mixing block for Wheeler-LM, the family's | |
| quantum-gravity mixer (structural analog of Quazimoto's Kuramoto block and Mycel's | |
| growth block). It is built from the Wheeler-DeWitt equation of canonical quantum | |
| gravity -- the "wavefunction of the universe": | |
| hbar^2 G_ijmn (d/dg_ij)(dPsi/dg_mn) + R sqrt(g) Psi = 0 (Hamiltonian constraint) | |
| (dPsi/dg_mn)_|m = 0 (momentum / diffeo constraint) | |
| What makes this different from every other mixer in the family: | |
| * NO TIME. It is not an evolution dPsi/dt = ...; it is a CONSTRAINT H Psi = 0. The | |
| state does not flow forward in an external time -- it must merely SATISFY the | |
| constraint. We realise this as an equilibrium/constraint block whose "time" is only | |
| the iteration index, plus an auxiliary loss <H^2> trained toward zero. | |
| * LORENTZIAN supermetric. The DeWitt supermetric G_ijmn is INDEFINITE -- one timelike | |
| direction (the volume / scale mode) and the rest spacelike. So the kinetic operator | |
| G d d is a WAVE operator, not a diffusion. We give the K minisuperspace modes a | |
| learnable indefinite metric G^{-1}, initialised diag(-1, +1, ..., +1). | |
| * CURVATURE POTENTIAL. The potential term is R sqrt(g) (scalar curvature), realised as | |
| a per-mode curvature weighting r that the token produces. | |
| * DIFFEO INVARIANCE. The momentum constraint says Psi depends only on the geometry, not | |
| the coordinates -> the readout uses only INVARIANTS of the mode vector (the timelike | |
| component and the norm of the spacelike part), never the raw coordinate values. | |
| Mechanics (per token): hidden -> (Psi, Pi) minisuperspace amplitude + conjugate momentum | |
| -> a few leapfrog steps of the Wheeler-DeWitt wave dynamics | |
| Pi <- Pi - dt * (r ⊙ Psi) # Pi_dot = -dU/dPsi with U = 1/2 r Psi^2 | |
| Psi <- Psi + dt * (G^{-1} @ Pi) # Psi_dot = G^{-1} Pi (indefinite -> wave) | |
| -> readout of diffeo-invariants -> hidden, behind a family gate. The Hamiltonian | |
| H = 1/2 Pi^T G^{-1} Pi + 1/2 sum_k r_k Psi_k^2 is exposed as last_constraint so the | |
| trunk can add <H^2> to the loss (the block is pressured onto the physical H=0 surface). | |
| """ | |
| 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 | |
| class WheelerDeWittBlock(nn.Module): | |
| """Wheeler-DeWitt wave mixer: hidden -> minisuperspace (Psi, Pi) -> leapfrog wave | |
| dynamics under a Lorentzian DeWitt supermetric with a curvature potential -> readout | |
| of diffeo-invariants. Drop-in for QuazimotoBlock.forward(x, ring_ctl, phase_seed).""" | |
| def __init__(self, cfg): | |
| super().__init__() | |
| self.cfg = cfg | |
| self.K = cfg.wdw_modes | |
| self.norm = RMSNorm(cfg.d_model) | |
| self.to_psi = nn.Linear(cfg.d_model, self.K) # initial wavefunction amplitude | |
| self.to_pi = nn.Linear(cfg.d_model, self.K) # conjugate momentum | |
| self.to_curv = nn.Linear(cfg.d_model, self.K) # per-mode curvature R*sqrt(g) | |
| # DeWitt supermetric G^{-1}: a LEARNABLE indefinite (Lorentzian) metric on the K | |
| # modes. Parameterised by a full matrix, symmetrised at use; a per-mode signature | |
| # (sigma, init [-1,+1,...,+1]) fixes the Lorentzian character -- mode 0 is timelike | |
| # (the volume / scale direction, the emergent "clock"), the rest spacelike. | |
| self.ginv_raw = nn.Parameter(torch.zeros(self.K, self.K)) | |
| sig = torch.ones(self.K); sig[0] = -1.0 # Lorentzian signature | |
| self.register_buffer("signature", sig) | |
| self.log_lapse = nn.Parameter(torch.zeros(1)) # lapse N: scales the intrinsic-time step | |
| # readout on diffeo-INVARIANTS of the final state: [Psi, Pi, timelike comp, ||space||]. | |
| # (never on raw coordinate values -> honours the momentum constraint's invariance.) | |
| hidden = cfg.readout_mult * cfg.d_model | |
| self.readout = nn.Sequential( | |
| nn.Linear(2 * self.K + 2, hidden), nn.GELU(), | |
| nn.Linear(hidden, cfg.d_model)) | |
| 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.psi_seed_gate = nn.Parameter(torch.zeros(1)) # zero-init -> no-op | |
| self.last_constraint = None # <H^2> for the trunk to collect | |
| def _supermetric(self): | |
| """The learnable Lorentzian DeWitt supermetric G^{-1} [K,K]. Symmetric, with the | |
| signature baked into the diagonal so it starts genuinely indefinite (timelike mode | |
| 0). ginv_raw learns the off-diagonal mode couplings on top of that.""" | |
| A = 0.5 * (self.ginv_raw + self.ginv_raw.t()) # symmetric part | |
| return A + torch.diag(self.signature) # Lorentzian at init (A=0) | |
| def forward(self, x, ring_ctl=None, phase_seed=None): | |
| cfg = self.cfg | |
| B, T, _ = x.shape | |
| h = self.norm(x) | |
| psi = self.to_psi(h) # [B,T,K] wavefunction amplitude | |
| # NOTE: the FractalSeed trait is a NO-OP in this project. The trained Wheeler-LM | |
| # checkpoint left every psi_seed_gate ~0 (seed on-vs-off flips 0/64 argmax tokens, | |
| # ~0.1% of logit scale), so the fractal table is effectively UNUSED here. It is | |
| # kept only for family parity / possible future continued training -- do not treat | |
| # it as an active input to the Wheeler-DeWitt dynamics. | |
| if phase_seed is not None: # optional fractal seed of Psi (gated ~0) | |
| seed = phase_seed[..., :self.K].view(B, T, self.K) | |
| psi = psi + torch.tanh(self.psi_seed_gate) * seed | |
| pi = self.to_pi(h) # [B,T,K] conjugate momentum | |
| r = self.to_curv(h) # [B,T,K] curvature potential R*sqrt(g) | |
| Ginv = self._supermetric() # [K,K] Lorentzian | |
| dt = cfg.wdw_dt * F.softplus(self.log_lapse).clamp(max=4.0) # lapse-scaled intrinsic time | |
| # leapfrog integration of the Wheeler-DeWitt wave equation on the K modes. Because | |
| # Ginv is INDEFINITE, the mode-0 (volume) direction is a genuine wave/"time" while | |
| # the spacelike modes oscillate -- a Klein-Gordon eq on minisuperspace. | |
| pi = pi - 0.5 * dt * (r * psi) # half kick | |
| for _ in range(cfg.wdw_steps): | |
| psi = psi + dt * torch.matmul(pi, Ginv.t()) # drift: Psi_dot = G^{-1} Pi | |
| last = (_ == cfg.wdw_steps - 1) | |
| pi = pi - (0.5 if last else 1.0) * dt * (r * psi) # kick: Pi_dot = -(R sqrt g) Psi | |
| # Hamiltonian constraint H = 1/2 Pi^T G^{-1} Pi + 1/2 sum_k r_k Psi_k^2. Physical | |
| # states satisfy H = 0; expose <H^2> so the trunk can pressure the block onto that | |
| # constraint surface (the closest thing to "solving" H Psi = 0 in a feed-forward net). | |
| kinetic = 0.5 * (torch.matmul(pi, Ginv.t()) * pi).sum(-1) # [B,T] | |
| potential = 0.5 * (r * psi * psi).sum(-1) # [B,T] | |
| H = kinetic + potential | |
| self.last_constraint = (H ** 2).mean() if self.training else None | |
| # diffeo-INVARIANT readout features: the timelike component and the norm of the | |
| # spacelike part are invariant under the residual coordinate symmetry. | |
| t_comp = psi[..., :1] # timelike (volume/scale) mode | |
| space_norm = psi[..., 1:].norm(dim=-1, keepdim=True) # ||spacelike Psi|| | |
| feat = torch.cat([psi, pi, t_comp, space_norm], dim=-1) | |
| out = self.drop(soft_clamp(self.readout(feat) * torch.tanh(self.gate), cfg.osc_bound)) | |
| rec = _viz.get_rec() | |
| if rec is not None and rec.enabled: # live-viz: emergent-time clock | |
| rec.log_ring([float(H[0, -1].abs())], [float(t_comp[0, -1, 0])], | |
| psi[0, -1].tolist()) | |
| rec.log_quaz_norm(out[0, -1].norm().item()) | |
| rec.flush_spec() | |
| return out | |