| 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)) |
| |
| 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) |
| |
| 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() |
| |
| 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 |
|
|
