from __future__ import annotations import math import torch from torch import nn from torch.nn import functional as F from .norms import RMSNorm, unit_rms class RelationalOperator(nn.Module): """Token-local asymmetric relations producing full per-unit SwiGLU FiLM.""" def __init__(self, d_model: int, observer_dim: int, dispersion_dim: int, ffn_hidden: int, *, modes: int = 12, eps: float = 1e-5, film_scale: float = 0.5, chain_scale: float = 0.5, include_depth: bool = True): super().__init__() if observer_dim != 48 and d_model == 768: raise ValueError("production observer_dim is frozen to 48") self.include_depth = include_depth self.modes = modes self.ffn_hidden = ffn_hidden self.film_scale = film_scale self.chain_scale = chain_scale self.evidence_names = (("f", "g", "qk", "v", "d") if include_depth else ("f", "g", "d")) names = ("b",) + self.evidence_names self.norms = nn.ModuleDict({name: RMSNorm(d_model, eps) for name in names}) self.obs = nn.ModuleDict({ name: nn.Linear(d_model, observer_dim, bias=False) for name in names }) self.dispersion_projection = nn.Linear(dispersion_dim, observer_dim, bias=False) self.lr_projection = nn.Linear(observer_dim, 2 * modes, bias=False) self.ratio_weight = nn.Parameter(torch.zeros(modes)) self.mode_weight = nn.Parameter(torch.empty(modes, 7)) self.mode_bias = nn.Parameter(torch.zeros(modes)) self.chain_raw = nn.Parameter(torch.zeros(modes, modes)) self.consistency_weight = nn.Parameter(torch.empty(modes, 4)) self.consistency_bias = nn.Parameter(torch.zeros(modes)) # Each relation mode is a learned basis field over every FFN unit. self.film = nn.Linear(modes, 4 * ffn_hidden, bias=False) self.steering = nn.Linear(modes, 1, bias=False) self.chain_raw._no_weight_decay = True self.steering.weight._no_weight_decay = True nn.init.normal_(self.mode_weight, mean=0.0, std=0.02) nn.init.normal_(self.consistency_weight, mean=0.0, std=0.02) nn.init.zeros_(self.film.weight) nn.init.zeros_(self.steering.weight) @staticmethod def _log_ratio(rms: torch.Tensor, base_rms: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: return torch.log(rms / base_rms).to(dtype) def directional_pair(self, present: torch.Tensor, evidence: torch.Tensor): """Return Present->evidence and evidence->Present mode products.""" p_left, p_right = self.lr_projection(present).chunk(2, dim=-1) e_left, e_right = self.lr_projection(evidence).chunk(2, dim=-1) return p_left * e_right, p_right * e_left def strict_lower_operator(self) -> torch.Tensor: return self.chain_scale * torch.tril(torch.tanh(self.chain_raw), diagonal=-1) @staticmethod def cross_evidence(alpha: torch.Tensor, left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: left_bar = (alpha * left).sum(dim=-2) right_bar = (alpha * right).sum(dim=-2) return left_bar * right_bar - ( alpha.square() * left * right ).sum(dim=-2) def forward(self, b, f, g, qk, v, d, dispersion): sources = {"b": b, "f": f, "g": g, "qk": qk, "v": v, "d": d} observations: dict[str, torch.Tensor] = {} rms: dict[str, torch.Tensor] = {} for name in ("b",) + self.evidence_names: normalized, source_rms = self.norms[name].forward_with_rms(sources[name]) observations[name] = self.obs[name](normalized) rms[name] = source_rms dispersion_observation = self.dispersion_projection( unit_rms(dispersion).to(b.dtype) ) evidence = torch.stack( [observations[name] for name in self.evidence_names] + [dispersion_observation], dim=-2, ) ratios = torch.stack( [self._log_ratio(rms[name], rms["b"], b.dtype) for name in self.evidence_names] + [torch.zeros_like(rms["b"], dtype=b.dtype)], dim=-2, ) p_left, p_right = self.lr_projection(observations["b"]).chunk(2, dim=-1) e_left, e_right = self.lr_projection(evidence).chunk(2, dim=-1) forward_pair = p_left.unsqueeze(-2) * e_right reverse_pair = p_right.unsqueeze(-2) * e_left scores = (forward_pair + reverse_pair) / math.sqrt(2.0) scores = scores + ratios * self.ratio_weight alpha = torch.softmax(scores, dim=-2) forward_mode = (alpha * forward_pair).sum(dim=-2) reverse_mode = (alpha * reverse_pair).sum(dim=-2) self_mode = (alpha * e_left * e_right).sum(dim=-2) # Exact sum_{i != j} alpha_i alpha_j L_i R_j. cross_mode = self.cross_evidence(alpha, e_left, e_right) ratio_mode = (alpha * ratios).sum(dim=-2) difference = forward_mode - reverse_mode primitive = torch.stack(( 0.5 * (forward_mode + reverse_mode), 0.5 * difference, forward_mode * reverse_mode, 0.5 * difference.square(), self_mode, cross_mode, ratio_mode, ), dim=-1) rho = F.silu((primitive * self.mode_weight).sum(dim=-1) + self.mode_bias) operator = self.strict_lower_operator() # F.linear(x, A) == x @ A.T; reverse reuses exactly A.T. y_forward = rho + F.linear(rho, operator) y_reverse = rho + F.linear(rho, operator.transpose(0, 1)) chain_difference = y_forward - y_reverse consistency = torch.stack(( 0.5 * (y_forward + y_reverse), 0.5 * chain_difference, y_forward * y_reverse, 0.5 * chain_difference.square(), ), dim=-1) relation = torch.tanh( rho + (consistency * self.consistency_weight).sum(dim=-1) + self.consistency_bias ) film = self.film(relation).unflatten(-1, (4, self.ffn_hidden)) return film.unbind(dim=-2), relation, alpha def modulate(self, gate, up, controls): sg, bg, su, bu = controls gate_rms = torch.sqrt( gate.float().square().mean(dim=-1, keepdim=True) + 1e-12 ).to(gate.dtype) up_rms = torch.sqrt( up.float().square().mean(dim=-1, keepdim=True) + 1e-12 ).to(up.dtype) gate = gate * (1 + self.film_scale * torch.tanh(sg)) + ( self.film_scale * gate_rms * torch.tanh(bg).to(gate.dtype) ) up = up * (1 + self.film_scale * torch.tanh(su)) + ( self.film_scale * up_rms * torch.tanh(bu).to(up.dtype) ) return gate, up def reorient(self, current: torch.Tensor, raw: torch.Tensor, relation: torch.Tensor): proposed = unit_rms(raw) present = unit_rms(current) gain = 1.0 + 0.5 * torch.tanh(self.steering(relation)) return unit_rms(present + gain * (proposed - present)), gain RelationController = RelationalOperator