import math import torch import torch.nn as nn import torch.nn.functional as F from config import DROPOUT MIXINGS = ("orthogonal", "softmax") class PersistentReasoner(nn.Module): """ Persistent multi-hypothesis refinement. Naming, honestly ---------------- Earlier versions called this "Persistent Hamiltonian Reasoning". It was not Hamiltonian in any sense a physicist would accept: `interaction = softmax(metric(H) H^T / sqrt(d))` is not Hermitian (W is unconstrained), and `H <- LayerNorm(H + g * dt * field)` is not unitary. Unitary evolution U = exp(-i H t) preserves norm and is invertible, so it *cannot* map two distinct states onto the same one. The implemented map was dissipative and provably contractive: it collapsed to a constant fixed point (mean pairwise cosine -> 1.000000, output independent of input to 1.9e-03). What v45 actually fixes ----------------------- The mixing across hypotheses -- the only part that can oversmooth -- is now an ORTHOGONAL operator, built by the Cayley transform of an antisymmetric coupling matrix: S = metric(H_norm) H_norm^T / sqrt(d) A = (S - S^T) / 2 (antisymmetric) X = dt * A / 2 U = (I - X)^{-1} (I + X) (orthogonal, exactly) Three consequences, all of them the point: * ||U H||_F = ||H||_F and U is invertible, so distinct hypotheses stay distinct. The v44 collapse is not merely discouraged, it is unreachable through the mixing path. * (I - X) is always invertible: the eigenvalues of a real antisymmetric matrix are purely imaginary, so 1 - i*lambda is never zero. No numerical guard is needed. * U's entries are SIGNED. v44's row-softmax made every mixing weight non-negative, so field_k was a convex combination of the hypotheses and cancellation was arithmetically impossible -- the module named "interference" could not interfere. Signed weights restore that possibility, and unlike v44's pre-audit unnormalized `tanh` (which was signed but unbounded, and is what caused the collapse in the first place) an orthogonal operator cannot blow up or contract. What is NOT claimed ------------------- The per-hypothesis refinement that follows the mixing -- the gated residual, the FFN, the LayerNorms -- is a standard residual block and is not norm-preserving. Only the hypothesis-mixing operator is orthogonal. This block as a whole is not unitary evolution and is not described as such anywhere. `mixing="softmax"` restores the v44 row-stochastic behaviour as a parameter-matched ablation arm. Anti-collapse machinery kept from v44 ------------------------------------- * Each step re-anchors to the original hypotheses H0 through a gated residual, so the map cannot forget its own input. * `dt`, `memory_decay`, `validator_weight` and `anchor_weight` are stored as LOGITS. They were once stored as post-sigmoid values by mistake: `nn.Parameter(tensor(0.10))` then `sigmoid(...)` gives an effective step size of 0.525, five times the intended 0.10, which directly accelerated the collapse. """ def __init__(self, dim, steps=3, keep_trajectory=False, mixing="orthogonal"): super().__init__() if mixing not in MIXINGS: raise ValueError(f"mixing must be one of {MIXINGS}, got {mixing!r}") self.steps = steps self.mixing = mixing # trajectory.append(H.detach().cpu()) forces a device->host sync on # every step of every forward pass. Only the visualization tools # consume it, so it is off during training/eval by default. self.keep_trajectory = keep_trajectory #################################################### # Coupling #################################################### self.metric = nn.Linear(dim, dim, bias=False) self.coupling = nn.Linear(dim, dim, bias=False) #################################################### # Interference network #################################################### self.interference = nn.Sequential( nn.Linear(dim, dim), nn.GELU(), nn.Linear(dim, dim), ) #################################################### # Adaptive gate #################################################### self.gate = nn.Sequential( nn.Linear(dim * 2, dim), nn.GELU(), nn.Linear(dim, dim), nn.Sigmoid(), ) #################################################### # Feedforward refinement #################################################### self.ffn = nn.Sequential( nn.Linear(dim, dim * 2), nn.GELU(), nn.Dropout(DROPOUT), nn.Linear(dim * 2, dim), ) self.norm1 = nn.LayerNorm(dim) self.norm2 = nn.LayerNorm(dim) self.norm3 = nn.LayerNorm(dim) #################################################### # Learnable evolution rate -- stored as a LOGIT. # logit(0.10) = -2.1972, so sigmoid(dt) starts at 0.10. #################################################### self.dt = nn.Parameter(torch.tensor(-2.1972)) #################################################### # Validator influence (logit; sigmoid -> 0.30) #################################################### self.validator_weight = nn.Parameter(torch.tensor(-0.8473)) #################################################### # Persistent memory (logit; sigmoid -> 0.80) #################################################### self.memory_decay = nn.Parameter(torch.tensor(1.3863)) #################################################### # Anchor back to the input hypotheses (logit; sigmoid -> 0.50) #################################################### self.anchor_weight = nn.Parameter(torch.tensor(0.0)) # ========================================================== # MIXING OPERATORS # ========================================================== def _orthogonal_mix(self, H, H_norm, mask, step_size): """ Returns (mixed, U). U is exactly orthogonal by construction. Computed in fp32 regardless of autocast: torch.linalg.solve is not implemented for half precision on several backends, and the matrices here are tiny (K <= 8), so there is nothing to gain from lower precision. """ dtype = H.dtype H32 = H.float() S = torch.einsum( "bkd,bjd->bkj", self.metric(H_norm).float(), H_norm.float() ) / math.sqrt(H.shape[-1]) A = 0.5 * (S - S.transpose(1, 2)) if mask is not None: m = mask.float() A = A * m.unsqueeze(1) * m.unsqueeze(2) X = 0.5 * step_size.float() * A eye = torch.eye(A.shape[1], device=A.device, dtype=A.dtype) eye = eye.unsqueeze(0).expand_as(A) U = torch.linalg.solve(eye - X, eye + X) mixed = torch.bmm(U, H32) return mixed.to(dtype), U.to(dtype) def _softmax_mix(self, H, H_norm, mask): """ v44 behaviour, kept as an ablation arm. Row-stochastic, hence a convex combination: field_k = sum_j a_kj H_j with every a_kj >= 0 can never cancel, and with an unmasked mean-like row it drives every hypothesis toward the same vector. """ logits = torch.einsum( "bkd,bjd->bkj", self.metric(H_norm), H_norm ) / math.sqrt(H.shape[-1]) if mask is not None: logits = logits.masked_fill(~mask.unsqueeze(1), float("-inf")) A = torch.softmax(logits, dim=-1) return torch.einsum("bkj,bjd->bkd", A, H), A # ========================================================== # FORWARD # ========================================================== def forward(self, H, potential=None, mask=None): """ H : (B, K, D) hypothesis embeddings potential: (B, K, D) optional validator guidance field mask : (B, K) bool, True for real hypotheses. Padded entries are excluded from the mixing so a zero-padded hypothesis can't drag the field toward the origin. """ H0 = H trajectory = [] interaction = None memory = torch.zeros_like(H) step_size = torch.sigmoid(self.dt) for _ in range(self.steps): H_norm = F.normalize(H, dim=-1) ################################################ # Mixing across hypotheses ################################################ if self.mixing == "orthogonal": field, interaction = self._orthogonal_mix( H, H_norm, mask, step_size ) else: field, interaction = self._softmax_mix(H, H_norm, mask) field = self.coupling(field) ################################################ # Interference network ################################################ field = field + self.interference(field) ################################################ # Persistent memory ################################################ decay = torch.sigmoid(self.memory_decay) memory = decay * memory + (1.0 - decay) * field field = field + memory ################################################ # Validator guidance ################################################ if potential is not None: weight = torch.sigmoid(self.validator_weight) field = field + weight * potential ################################################ # Gated evolution ################################################ gate = self.gate(torch.cat([H, field], dim=-1)) H = self.norm1(H + gate * step_size * field) ################################################ # Refinement ################################################ H = self.norm2(H + self.ffn(H)) ################################################ # Re-anchor to the original hypotheses so the map # cannot converge to an input-independent point. ################################################ H = self.norm3(H + torch.sigmoid(self.anchor_weight) * H0) if mask is not None: H = H * mask.unsqueeze(-1).to(H.dtype) if self.keep_trajectory: trajectory.append(H.detach().cpu()) return H, trajectory, interaction