File size: 7,382 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 | """Complete Toroidal Fractal Intelligence pipeline."""
from __future__ import annotations
import torch
import torch.nn as nn
from .encoder import ToroidalEncoder
from .state import FractalSuperpositionState
from .dynamics import RK4DynamicsEngine, ToroidalDynamics
from .interaction import ToroidalInteraction
from .aggregation import AggregationEngine
from .abstraction import AbstractionEngine
from .consolidation import ConsolidationEngine
from .production import ProductionHead
from .atom import ToroidalAtomCollection
class ToroidalFractalIntelligence(nn.Module):
def __init__(self, vocab_size=32000, d_model=256, n_modes=256, n_atoms_max=1024):
super().__init__()
self.encoder = ToroidalEncoder(vocab_size, d_model, n_atoms_max=n_atoms_max)
self.state = FractalSuperpositionState(n_modes, d_model)
self.dynamics = RK4DynamicsEngine(ToroidalDynamics(d_model, n_modes), dt=0.01, n_steps=4)
self.interaction = ToroidalInteraction(d_model, n_modes)
self.aggregation = AggregationEngine(d_model)
self.abstraction = AbstractionEngine(d_model)
self.consolidation = ConsolidationEngine(d_model)
self.production = ProductionHead(d_model, vocab_size)
self.atoms = ToroidalAtomCollection()
self.energy_history = []
self.training_loss = 0.0
self.consolidation_count = 0
self.abstraction_count = 0
def forward(self, token_ids, context_ids=None):
if token_ids.dim() > 1:
token_ids = token_ids[:, -1]
new_atom, operation, confidence = self.encoder(token_ids, context_ids, self.atoms)
self.state.add_atom_contribution(new_atom.r, new_atom.phi, new_atom.omega, new_atom.E, new_atom.kappa)
self.atoms.add(new_atom)
# The committed state is recurrent memory; use a snapshot so committing
# the next state cannot invalidate the autograd graph for this tick.
alpha = self.state.get_field().detach().clone()
alpha_new = self.dynamics.evolve(alpha, input_token=new_atom.r)
interaction_energy = self.interaction.field_interaction(alpha_new)
# Explicit toroidal interaction feeds the resulting state.
left, right = torch.roll(alpha_new, 1, 0), torch.roll(alpha_new, -1, 0)
alpha_new = alpha_new + 0.1 * self.dynamics.dynamics.phase_sync * (left + right - 2 * alpha_new)
aggregates = []
if len(self.atoms) > 10:
aggregates, _ = self.aggregation.aggregate(self.atoms.r, self.atoms.phi, self.atoms.omega, self.atoms.E, self.atoms.kappa, self.atoms.M, self.atoms.tau, self.atoms.rho)
abstractions = []
for agg in aggregates:
idx = self.abstraction.create_or_update_abstraction([agg])
if idx >= 0:
abstractions.append({"id": idx, "pattern": self.abstraction.abstraction_memory[idx].detach()})
self.abstraction_count += 1
self.abstraction.step_age()
persistent_state = None
if len(self.energy_history) > 5:
E = self.atoms.E.mean(dim=0)
stability = self.consolidation.compute_stability(E, self.energy_history)
persistent_state, mask = self.consolidation.consolidate(alpha_new, E.unsqueeze(0), stability.unsqueeze(0))
if mask.any(): self.consolidation_count += int(mask.sum().item())
self.energy_history.append(self.atoms.E.mean(dim=0).detach())
self.energy_history = self.energy_history[-20:]
logits, conf = self.production.produce(alpha_new, persistent_state, abstractions)
logits = logits.expand(token_ids.shape[0], -1)
with torch.no_grad():
self.state.alpha.copy_(alpha_new.detach())
self.state.t.add_(self.dynamics.dt)
return {"logits": logits, "field": alpha_new, "confidence": conf, "interaction_energy": interaction_energy, "aggregates": aggregates, "abstractions": abstractions, "persistent_state": persistent_state, "operation": operation}
def train_step(self, token_ids, target_ids, optimizer, context_ids=None):
optimizer.zero_grad(set_to_none=True)
output = self.forward(token_ids, context_ids)
if target_ids.dim() > 1: target_ids = target_ids[:, -1]
loss = nn.CrossEntropyLoss()(output["logits"], target_ids)
alpha = output["field"]
total_loss = loss + 0.01 * alpha.abs().mean()
total_loss.backward()
torch.nn.utils.clip_grad_norm_(self.parameters(), 1.0)
optimizer.step()
self.training_loss = loss.item()
return {"loss": loss.item(), "total_loss": total_loss.item(), "n_atoms": len(self.atoms), "n_aggregates": len(output["aggregates"]), "n_abstractions": len(output["abstractions"]), "consolidation_count": self.consolidation_count}
def generate(self, prompt_ids, max_length=100, temperature=1.0, top_k=50):
generated = prompt_ids.clone()
for _ in range(max_length):
out = self.forward(generated[-1:].unsqueeze(0))
token_id, _, _ = self.production.sample(self.state.get_field(), out["persistent_state"], out["abstractions"], temperature, top_k)
generated = torch.cat([generated, torch.tensor([token_id], device=generated.device)])
if token_id == 2: break
return generated
def save(self, path, extra_state=None):
checkpoint = {"encoder": self.encoder.state_dict(), "state": self.state.get_state_dict(), "dynamics": self.dynamics.state_dict(), "interaction": self.interaction.state_dict(), "aggregation": self.aggregation.state_dict(), "abstraction": self.abstraction.state_dict(), "consolidation": self.consolidation.state_dict(), "production": self.production.state_dict(), "atoms": self.atoms.state_dict(), "energy_history": self.energy_history, "training_loss": self.training_loss, "consolidation_count": self.consolidation_count, "abstraction_count": self.abstraction_count, "config": {"vocab_size": self.encoder.vocab_size, "d_model": self.encoder.d_model, "n_modes": self.state.n_modes, "n_atoms_max": self.encoder.n_atoms_max}}
if extra_state is not None:
checkpoint["training"] = extra_state
torch.save(checkpoint, path)
def load(self, path):
checkpoint = torch.load(path, map_location="cpu", weights_only=False)
for name in ("encoder", "dynamics", "interaction", "aggregation", "abstraction", "consolidation", "production"): getattr(self, name).load_state_dict(checkpoint[name])
self.state.load_state_dict(checkpoint["state"]); self.atoms.load_state_dict(checkpoint["atoms"])
self.energy_history = checkpoint.get("energy_history", []); self.training_loss = checkpoint.get("training_loss", 0.0); self.consolidation_count = checkpoint.get("consolidation_count", 0); self.abstraction_count = checkpoint.get("abstraction_count", 0)
return checkpoint.get("training")
def get_shared_params_summary(self):
return {"coupling_scale": self.dynamics.dynamics.coupling_scale.item(), "energy_decay": self.dynamics.dynamics.energy_decay.item(), "phase_sync": self.dynamics.dynamics.phase_sync.item(), "n_atoms": len(self.atoms), "n_modes": self.state.n_modes, "d_model": self.encoder.d_model}
def create_model(vocab_size=32000, d_model=256, n_modes=256, n_atoms_max=1024):
return ToroidalFractalIntelligence(vocab_size, d_model, n_modes, n_atoms_max)
|