File size: 5,872 Bytes
b9301f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """Explicit toroidal field dynamics and RK4 integration.
The dynamics layer deliberately contains no attention/Transformer block and no
learned feature-mixing MLP. The state evolves through shared physical-style
rules: toroidal neighbour coupling, phase rotation, input driving, and energy
decay. RK4 is only the numerical integrator for this continuous system.
"""
from __future__ import annotations
import torch
import torch.nn as nn
class ToroidalDynamics(nn.Module):
"""Shared-rule dynamics for the spectral toroidal field.
``alpha[m]`` is the amplitude/state carried by toroidal mode ``m``.
Mode phase is fixed by its position on the torus; the incoming atom is
injected as an external drive. All trainable dynamics parameters are
shared globally rather than being per-mode feature transforms.
"""
def __init__(self, d_model: int = 256, n_modes: int = 256) -> None:
super().__init__()
self.d_model = d_model
self.n_modes = n_modes
self.coupling_scale = nn.Parameter(torch.tensor(1.0))
self.energy_decay = nn.Parameter(torch.tensor(0.99))
self.phase_sync = nn.Parameter(torch.tensor(0.5))
self.register_buffer(
"mode_phase",
torch.arange(n_modes, dtype=torch.float32) * (2.0 * torch.pi / n_modes),
)
def forward(
self,
alpha: torch.Tensor,
input_token: torch.Tensor | None = None,
context: torch.Tensor | None = None,
) -> torch.Tensor:
if alpha.shape[-2:] != (self.n_modes, self.d_model):
raise ValueError(
f"expected alpha [..., {self.n_modes}, {self.d_model}], got {tuple(alpha.shape)}"
)
# Toroidal nearest-neighbour coupling. Periodic roll closes the ring.
left = torch.roll(alpha, shifts=1, dims=-2)
right = torch.roll(alpha, shifts=-1, dims=-2)
laplacian = left + right - 2.0 * alpha
coupling = self.phase_sync * laplacian
# Continuous phase rotation of each mode. This is the field analogue
# of oscillator evolution; it is not token attention.
phase = self.mode_phase.to(device=alpha.device, dtype=alpha.dtype)
phase = phase.view(*([1] * (alpha.ndim - 2)), self.n_modes, 1)
rotated = torch.sin(phase) * alpha
# External token/context drive. Mode-modulated (choice A): each mode
# gets a content×phase Hadamard of the token — not the same broadcast
# vector on every mode (mean-broadcast made distinct texts collinear).
drive = torch.zeros_like(alpha)
dim_phase = torch.arange(
self.d_model, device=alpha.device, dtype=alpha.dtype
).view(*([1] * (alpha.ndim - 1)), self.d_model)
if input_token is not None:
token = input_token
while token.ndim > 1:
token = token.mean(dim=0)
tok = token.reshape(*([1] * (alpha.ndim - 2)), 1, self.d_model)
# drive_m,d = token_d * cos(mode_phase_m + 0.5 * token_d * d)
drive = tok * torch.cos(phase + 0.5 * tok * dim_phase)
drive = drive.expand(*alpha.shape[:-2], self.n_modes, self.d_model)
drive = self.coupling_scale * drive
if context is not None:
ctx = context
while ctx.ndim > 1:
ctx = ctx.mean(dim=0)
ctx = ctx.reshape(*([1] * (alpha.ndim - 2)), 1, self.d_model)
ctx = ctx * torch.cos(phase + 0.5 * ctx * dim_phase)
ctx = ctx.expand(*alpha.shape[:-2], self.n_modes, self.d_model)
drive = drive + 0.5 * self.coupling_scale * ctx
decay = (self.energy_decay - 1.0) * alpha
return coupling + rotated + drive + decay
def _mode_freqs(self, device: torch.device) -> torch.Tensor:
return torch.arange(self.n_modes, device=device, dtype=torch.float32) / self.n_modes * 2 * torch.pi
def _compute_phase_laplacian(self, alpha: torch.Tensor) -> torch.Tensor:
left = torch.roll(alpha, shifts=1, dims=-2)
right = torch.roll(alpha, shifts=-1, dims=-2)
return torch.sin(left - alpha) + torch.sin(right - alpha)
def get_shared_params(self) -> dict:
return {
"coupling_scale": self.coupling_scale.detach().cpu(),
"energy_decay": self.energy_decay.detach().cpu(),
"phase_sync": self.phase_sync.detach().cpu(),
}
def load_shared_params(self, params: dict) -> None:
for key in ("coupling_scale", "energy_decay", "phase_sync"):
if key in params:
getattr(self, key).data.copy_(params[key].to(getattr(self, key).device))
class RK4DynamicsEngine(nn.Module):
"""RK4 numerical integration of the toroidal dynamics."""
def __init__(self, dynamics: ToroidalDynamics, dt: float = 0.01, n_steps: int = 4) -> None:
super().__init__()
self.dynamics = dynamics
self.dt = dt
self.n_steps = n_steps
def evolve(
self,
alpha: torch.Tensor,
input_token: torch.Tensor | None = None,
context: torch.Tensor | None = None,
) -> torch.Tensor:
def f(state: torch.Tensor) -> torch.Tensor:
return self.dynamics(state, input_token, context)
k1 = f(alpha)
k2 = f(alpha + self.dt * 0.5 * k1)
k3 = f(alpha + self.dt * 0.5 * k2)
k4 = f(alpha + self.dt * k3)
return alpha + (self.dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
def evolve_multi(
self,
alpha: torch.Tensor,
input_token: torch.Tensor | None = None,
context: torch.Tensor | None = None,
n_steps: int | None = None,
) -> torch.Tensor:
current = alpha
for _ in range(n_steps or self.n_steps):
current = self.evolve(current, input_token, context)
return current
|