File size: 2,537 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 | """Token + context -> structured toroidal atom."""
from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
from .atom import ToroidalAtom
class ToroidalEncoder(nn.Module):
"""Shared encoder that creates/modifies computational matter."""
def __init__(self, vocab_size=32000, d_model=256, d_hidden=512, n_atoms_max=1024, tau_eps=1e-6):
super().__init__()
self.vocab_size, self.d_model, self.n_atoms_max, self.tau_eps = vocab_size, d_model, n_atoms_max, tau_eps
self.token_embedding = nn.Embedding(vocab_size, d_model)
self.mapping = nn.Sequential(nn.Linear(d_model, d_hidden), nn.GELU(), nn.Linear(d_hidden, 8 * d_model))
self.op_gate = nn.Linear(d_model, 7)
self.theta_shared = nn.ParameterDict({
"coupling_scale": nn.Parameter(torch.tensor(1.0)),
"energy_decay": nn.Parameter(torch.tensor(0.99)),
"phase_sync_strength": nn.Parameter(torch.tensor(0.5)),
})
def forward(self, token_ids, context_ids=None, existing_atoms=None):
tok = self.token_embedding(token_ids)
if context_ids is not None:
ctx = self.token_embedding(context_ids).mean(dim=-2) if context_ids.dim() > 1 else self.token_embedding(context_ids)
tok = tok + 0.25 * ctx
raw = self.mapping(tok).view(tok.shape[0], 8, self.d_model)
r = F.normalize(raw[:, 0], dim=-1)
phi = raw[:, 1]
omega = F.softplus(raw[:, 2]) + 1e-3
E = torch.sigmoid(raw[:, 3])
kappa = F.softplus(raw[:, 4])
M = F.normalize(raw[:, 5], dim=-1)
tau = F.softplus(raw[:, 6]) + self.tau_eps
# Hierarchical level is structural metadata, not random noise.
base_rho = len(existing_atoms) if existing_atoms is not None else 0
rho = torch.full((tok.shape[0], 1), min(base_rho // 64, 4), dtype=tok.dtype, device=tok.device)
new_atom = ToroidalAtom(r=r, phi=phi, omega=omega, E=E, kappa=kappa, M=M, tau=tau, rho=rho)
op_logits = self.op_gate(tok)
operation = op_logits.argmax(dim=-1)
confidence = op_logits.softmax(dim=-1).max(dim=-1).values
return new_atom, operation, confidence
def get_shared_params(self):
return {k: v.detach().clone() for k, v in self.theta_shared.items()}
def load_shared_params(self, params):
with torch.no_grad():
for k, v in params.items():
if k in self.theta_shared:
self.theta_shared[k].copy_(v)
|