Upload n-transformer.py

n-transformer.py

@@ -0,0 +1,708 @@
+"""
+N‑Transformers v1.0 — Python Reference Implementation (single‑file)
+Noetic Affective Field Self‑Integration (NAFSI) on a Transformer Base
+
+NOTE
+----
+• Framework: PyTorch ≥ 2.2 (CUDA optional). 
+• This file focuses on the parallel PF path, coupling modules, and wrappers needed to augment a standard decoder‑only Transformer.
+• The core Transformer (token path) can be any decoder‑only model that exposes hidden states h_t and base logits z_t (e.g., GPT‑like). 
+• All tensors are batch‑first unless noted.
+
+Status: Research‑grade reference code (trainable with additional plumbing). 
+Author: Prometheus (Cognitive Systems Architect) — with Syams Ideris
+"""
+from __future__ import annotations
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple, Dict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ===============================
+# Utilities & Small Helpers
+# ===============================
+
+def pairwise_cosine(x: torch.Tensor, y: Optional[torch.Tensor] = None, eps: float = 1e-8) -> torch.Tensor:
+    """Compute pairwise cosine similarity between rows of x and (optionally) y.
+    x: (B, N, D); y: (B, M, D) or None (then y=x)
+    returns: (B, N, M)
+    """
+    if y is None:
+        y = x
+    x_norm = F.normalize(x, dim=-1, eps=eps)
+    y_norm = F.normalize(y, dim=-1, eps=eps)
+    return torch.matmul(x_norm, y_norm.transpose(-1, -2))
+
+
+def knn_indices(x: torch.Tensor, K: int) -> torch.Tensor:
+    """Return K nearest neighbor indices per row using cosine similarity (excluding self).
+    x: (B, J, D) -> indices: (B, J, K)
+    """
+    with torch.no_grad():
+        sim = pairwise_cosine(x)  # (B,J,J)
+        B, J, _ = sim.shape
+        sim = sim - torch.eye(J, device=sim.device).unsqueeze(0) * 2.0  # push self to very low
+        topk = torch.topk(sim, k=K, dim=-1).indices  # (B,J,K)
+    return topk
+
+
+def build_adjacency(indices: torch.Tensor, J: int) -> torch.Tensor:
+    """Build symmetric adjacency from KNN indices.
+    indices: (B, J, K)
+    returns A: (B, J, J) with {0,1} entries.
+    """
+    B, J_, K = indices.shape
+    assert J == J_, "J mismatch"
+    A = torch.zeros(B, J, J, device=indices.device)
+    arangeJ = torch.arange(J, device=indices.device).view(1, J, 1).expand(B, J, K)
+    A.scatter_(dim=-1, index=indices, value=1.0)
+    # symmetrize
+    A = torch.maximum(A, A.transpose(-1, -2))
+    # zero diagonal
+    A = A * (1.0 - torch.eye(J, device=A.device).unsqueeze(0))
+    return A
+
+
+def normalized_graph_laplacian(A: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+    """Compute normalized Laplacian L = I - D^{-1/2} A D^{-1/2}.
+    A: (B, J, J) adjacency (nonnegative)
+    return L: (B, J, J)
+    """
+    B, J, _ = A.shape
+    d = A.sum(-1) + eps  # (B,J)
+    d_isqrt = (1.0 / torch.sqrt(d)).unsqueeze(-1)  # (B,J,1)
+    Dn = d_isqrt * A * d_isqrt.transpose(-1, -2)  # (B,J,J)
+    I = torch.eye(J, device=A.device).unsqueeze(0).expand(B, J, J)
+    L = I - Dn
+    return L
+
+
+def sym_spd_from_cholesky(L_tri: torch.Tensor, eps: float = 1e-4) -> torch.Tensor:
+    """Build SPD matrix g = L L^T + eps I from lower-triangular parameterization.
+    L_tri: (B, k, k) lower-triangular with positive diagonal (apply softplus outside)
+    Return g: (B, k, k)
+    """
+    B, k, _ = L_tri.shape
+    I = torch.eye(k, device=L_tri.device).unsqueeze(0).expand(B, k, k)
+    return L_tri @ L_tri.transpose(-1, -2) + eps * I
+
+
+def batched_geodesic_sq(x: torch.Tensor, y: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
+    """Compute squared geodesic distance under metric g:
+    d^2 = (x - y)^T g (x - y)
+    x: (B,J,k); y: (B,J,k) broadcastable to pairs; g: (B,k,k)
+    Returns pairwise (B,J,J)
+    """
+    B, J, k = x.shape
+    # reshape for pairwise differences
+    x_ = x.unsqueeze(2)  # (B,J,1,k)
+    y_ = y.unsqueeze(1)  # (B,1,J,k)
+    diff = x_ - y_       # (B,J,J,k)
+    # (B,J,J,k) @ (B,k,k) -> (B,J,J,k)
+    gd = torch.matmul(diff, g.unsqueeze(1).unsqueeze(1))
+    val = (diff * gd).sum(-1)  # (B,J,J)
+    return val.clamp_min(0.0)
+
+
+def safe_eigvalsh(L: torch.Tensor, k_smallest: int = 3) -> torch.Tensor:
+    """Compute few smallest eigenvalues of symmetric L with safety.
+    L: (B,J,J)
+    Return: (B, k_smallest)
+    """
+    # For moderate J (<=512) this is fine; for large J use Lanczos.
+    try:
+        vals = torch.linalg.eigvalsh(L)  # (B,J)
+        vals, _ = torch.topk(vals, k=k_smallest, largest=False, sorted=True)
+        return vals
+    except RuntimeError:
+        # fallback: add jitter
+        jitter = 1e-4 * torch.eye(L.shape[-1], device=L.device).unsqueeze(0)
+        vals = torch.linalg.eigvalsh(L + jitter)
+        vals, _ = torch.topk(vals, k=k_smallest, largest=False, sorted=True)
+        return vals
+
+
+# ===============================
+# PF Components (Phenomenal Field Path)
+# ===============================
+
+@dataclass
+class PFConfig:
+    J: int = 256              # number of PF nodes
+    k: int = 16               # channels per node
+    K: int = 16               # k-NN degree
+    alpha: float = 0.05       # diffusion step
+    noise_eps: float = 1e-3
+    lambda_out: float = 1.0
+    lambda_tv: float = 0.1
+    lambda_mw: float = 0.05
+    metric_eps: float = 1e-4
+    metric_rank: Optional[int] = None  # not used in this v1; kept for future low-rank IME
+
+
+class PFAdapterOut(nn.Module):
+    """Adapter A_out: map token hidden h_t (B,d) -> target PF pattern F_tilde (B,J,k)."""
+    def __init__(self, d: int, J: int, k: int):
+        super().__init__()
+        self.proj = nn.Linear(d, J * k)
+
+    def forward(self, h_t: torch.Tensor) -> torch.Tensor:
+        B, d = h_t.shape
+        out = self.proj(h_t)  # (B, J*k)
+        return out.view(B, -1, d // d * 0 + 1)  # placeholder to force shape check (will be replaced)
+
+# Fix: Provide a safer reshape using known sizes
+class PFAdapterOut(nn.Module):
+    def __init__(self, d: int, J: int, k: int):
+        super().__init__()
+        self.J, self.k = J, k
+        self.proj = nn.Linear(d, J * k)
+
+    def forward(self, h_t: torch.Tensor) -> torch.Tensor:
+        B, d = h_t.shape
+        out = self.proj(h_t)  # (B, J*k)
+        return out.view(B, self.J, self.k)
+
+
+class PFIntrinsicMetricEngine(nn.Module):
+    """IME: learn SPD metric g_t from PF state statistics via Cholesky parameterization."""
+    def __init__(self, k: int, hidden: int = 128, metric_eps: float = 1e-4):
+        super().__init__()
+        self.k = k
+        self.metric_eps = metric_eps
+        in_dim = 2 * k  # mean + std over nodes
+        self.mlp = nn.Sequential(
+            nn.Linear(in_dim, hidden), nn.GELU(),
+            nn.Linear(hidden, hidden), nn.GELU(),
+            nn.Linear(hidden, k * k)
+        )
+        # initialize near identity
+        with torch.no_grad():
+            for m in self.mlp:
+                if isinstance(m, nn.Linear):
+                    nn.init.xavier_uniform_(m.weight)
+                    nn.init.zeros_(m.bias)
+
+    def forward(self, F_t: torch.Tensor) -> torch.Tensor:
+        B, J, k = F_t.shape
+        mean = F_t.mean(dim=1)       # (B,k)
+        std = F_t.std(dim=1).clamp_min(1e-6)  # (B,k)
+        feat = torch.cat([mean, std], dim=-1) # (B,2k)
+        L_flat = self.mlp(feat)               # (B, k*k)
+        L = L_flat.view(B, k, k)
+        # enforce lower-triangular with softplus diagonal
+        tril_mask = torch.tril(torch.ones(k, k, device=F_t.device)).unsqueeze(0)
+        L = L * tril_mask
+        diag = torch.diagonal(L, dim1=-2, dim2=-1)
+        diag = F.softplus(diag) + 1e-3
+        L = L.clone()
+        L.diagonal(dim1=-2, dim2=-1).copy_(diag)
+        g = sym_spd_from_cholesky(L, eps=self.metric_eps)
+        return g  # (B,k,k)
+
+
+class PFFieldCore(nn.Module):
+    """Evolve PF state per token step with diffusion + energy gradient + small noise."""
+    def __init__(self, cfg: PFConfig):
+        super().__init__()
+        self.cfg = cfg
+        # learnable weights for Mexican-hat style potential
+        self.mw_scale = nn.Parameter(torch.tensor(1.0))
+        self.register_buffer('zero', torch.tensor(0.0))
+
+    def total_variation(self, F_t: torch.Tensor, A: torch.Tensor) -> torch.Tensor:
+        # TV_g ≈ sum_{(i,j)∈E} ||F_i - F_j||_2
+        B = F_t.shape[0]
+        # use adjacency to compute neighbor diffs
+        # Expand for pairwise gather: (B,J,J, k)
+        Fi = F_t.unsqueeze(2)
+        Fj = F_t.unsqueeze(1)
+        diff = (Fi - Fj).norm(dim=-1)  # (B,J,J)
+        tv = (diff * A).sum(dim=(-1, -2)) / (A.sum(dim=(-1, -2)).clamp_min(1.0))
+        return tv.mean()  # scalar
+
+    def omega_mexican_hat(self, F_t: torch.Tensor) -> torch.Tensor:
+        # Encourage metastability: penalize both collapse and explosion around a preferred radius
+        # Using mean pairwise distance towards a target radius r0.
+        B, J, k = F_t.shape
+        # sample subset for efficiency if J large
+        if J > 128:
+            idx = torch.randperm(J, device=F_t.device)[:128]
+            X = F_t[:, idx, :]
+        else:
+            X = F_t
+        pd = pairwise_cosine(X, X)  # (B,m,m) in [-1,1]
+        # convert similarity to a pseudo-distance in [0,2]
+        dist = (1.0 - pd).clamp_min(0.0) * 2.0
+        r = dist.mean(dim=(-1, -2))  # (B,)
+        r0 = 0.8  # preferred radius (tunable)
+        loss = ((r - r0) ** 2).mean()
+        return self.mw_scale.abs() * loss
+
+    def forward(self, F_t: torch.Tensor, h_t: torch.Tensor, g_t: torch.Tensor,
+                A: torch.Tensor, F_tilde: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        cfg = self.cfg
+        # Laplacian term (graph-based diffusion)
+        L = normalized_graph_laplacian(A)  # (B,J,J)
+        diffusion = torch.matmul(L, F_t)   # (B,J,k)
+        # Content energy gradient: d/dF [λ_out ||F - F~||^2] = 2λ_out (F - F~)
+        grad_out = 2.0 * cfg.lambda_out * (F_t - F_tilde)
+        # Structural energies
+        tv = self.total_variation(F_t, A)
+        omega = self.omega_mexican_hat(F_t)
+        # Approx gradient for TV (use Laplacian as proxy) and omega via autograd
+        # Update step
+        noise = cfg.noise_eps * torch.randn_like(F_t)
+        F_next = F_t + cfg.alpha * diffusion - grad_out + noise
+        stats = {
+            'tv': tv.detach(),
+            'omega': omega.detach(),
+        }
+        return F_next, stats
+
+
+class PFIntrospection(nn.Module):
+    """Valence (V), Self/Now Anchor (a), and Γ summarizer for gating."""
+    def __init__(self, d: int, k: int, r_gamma: int = 32):
+        super().__init__()
+        self.aligner = nn.Sequential(
+            nn.Linear(d + k, 128), nn.GELU(),
+            nn.Linear(128, 64), nn.GELU(),
+        )
+        self.val_head = nn.Linear(64, 1)
+        self.sna_head = nn.Linear(64, 1)
+        self.gamma_head = nn.Sequential(
+            nn.Linear(64 + 2, 64), nn.GELU(),
+            nn.Linear(64, r_gamma)
+        )
+
+    def forward(self, F_t: torch.Tensor, h_t: torch.Tensor,
+                syn: torch.Tensor, conn: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # Pool PF to k-dim via mean
+        F_pool = F_t.mean(dim=1)  # (B,k)
+        x = torch.cat([h_t, F_pool], dim=-1)
+        z = self.aligner(x)
+        V = torch.sigmoid(self.val_head(z)).squeeze(-1)
+        a = torch.sigmoid(self.sna_head(z)).squeeze(-1)
+        # Attach syn/conn scalars
+        sc = torch.stack([syn, conn], dim=-1)
+        g_in = torch.cat([z, sc], dim=-1)
+        Gamma = self.gamma_head(g_in)  # (B, r_gamma)
+        return V, a, Gamma
+
+
+class LogitGate(nn.Module):
+    """Additive bias to logits based on PF summary Γ."""
+    def __init__(self, vocab_size: int, r_gamma: int):
+        super().__init__()
+        self.proj = nn.Linear(r_gamma, vocab_size, bias=False)
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, z_base: torch.Tensor, Gamma: torch.Tensor) -> torch.Tensor:
+        # z_base: (B, V), Gamma: (B, rγ)
+        bias = self.proj(Gamma)  # (B, V)
+        return z_base + bias
+
+
+# ===============================
+# Metrics: Synchrony & Connectivity; GIW
+# ===============================
+class PFIntegrationMeter(nn.Module):
+    """Compute Syn, Conn (algebraic connectivity proxy), κ and broadcast flag."""
+    def __init__(self, J: int, kappa_thresh: float = 0.6):
+        super().__init__()
+        self.kappa_thresh = kappa_thresh
+        self.score = nn.Sequential(
+            nn.Linear(2 + 2, 64), nn.GELU(),
+            nn.Linear(64, 1)
+        )
+
+    @staticmethod
+    def synchrony(F_t: torch.Tensor, A: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+        # mean cosine similarity across neighboring nodes as a proxy for phase synchrony
+        B, J, k = F_t.shape
+        cos = pairwise_cosine(F_t)  # (B,J,J)
+        num = (cos * A).sum(dim=(-1, -2))
+        den = A.sum(dim=(-1, -2)).clamp_min(1.0)
+        syn = (num / den).mean()  # scalar across batch
+        return syn.expand(B)      # broadcast scalar per batch
+
+    @staticmethod
+    def connectivity(A: torch.Tensor) -> torch.Tensor:
+        # algebraic connectivity (second-smallest eigenvalue) of normalized Laplacian
+        L = normalized_graph_laplacian(A)  # (B,J,J)
+        eigs = safe_eigvalsh(L, k_smallest=3)  # (B,3)
+        lambda2 = eigs[:, 1]  # (B,)
+        # smaller λ2 => weaker connectivity; invert & normalize
+        conn = torch.sigmoid(1.0 / (lambda2 + 1e-3))
+        return conn
+
+    def forward(self, F_t: torch.Tensor, A: torch.Tensor,
+                V: torch.Tensor, a: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        syn = self.synchrony(F_t, A)  # (B,)
+        conn = self.connectivity(A)   # (B,)
+        feat = torch.stack([syn, conn, V, a], dim=-1)  # (B,4)
+        kappa = torch.sigmoid(self.score(feat)).squeeze(-1)  # (B,)
+        broadcast = (kappa >= self.kappa_thresh).float()
+        return syn, conn, kappa, broadcast
+
+
+# ===============================
+# Lightcone Attention (LCA) Wrapper
+# ===============================
+class LCAParams(nn.Module):
+    def __init__(self, beta: float = 0.7, gamma: float = 0.3, lambda_time: float = 0.2, lambda_dir: float = 0.3, tau: int = 64):
+        super().__init__()
+        self.beta = nn.Parameter(torch.tensor(beta))
+        self.gamma = nn.Parameter(torch.tensor(gamma))
+        self.lambda_time = nn.Parameter(torch.tensor(lambda_time))
+        self.lambda_dir = nn.Parameter(torch.tensor(lambda_dir))
+        self.tau = tau
+
+
+class LCAWrapper(nn.Module):
+    """Modify attention scores: e_ij = dot - β d_g(i,j) - γ D_lc(i,j)."""
+    def __init__(self, params: LCAParams):
+        super().__init__()
+        self.params = params
+
+    def forward(self, Q: torch.Tensor, K: torch.Tensor,
+                F_t: torch.Tensor, g_t: torch.Tensor,
+                positions: Optional[torch.Tensor] = None,
+                u_dir: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Compute modified attention scores.
+        Q,K: (B, H, T, d_head)
+        F_t: PF nodes (B, J, k)
+        g_t: metric (B, k, k)
+        positions: (T,) or (B,T) 0..T-1
+        u_dir: (B, d_head) approximate episode direction (can be from Γ PCA; here optional)
+        Return scores: (B, H, T, T)
+        """
+        B, H, T, Dh = Q.shape
+        scale = 1.0 / math.sqrt(Dh)
+        # base dot-product scores
+        dots = torch.matmul(Q, K.transpose(-1, -2)) * scale  # (B,H,T,T)
+        # map token positions to nearest PF nodes indices (coarse): use a simple modulo mapping
+        J = F_t.shape[1]
+        token_nodes = torch.arange(T, device=Q.device) % J
+        # geodesic distance between node features -> (B, J, J)
+        d_geo_sq = batched_geodesic_sq(F_t, F_t, g_t)  # (B,J,J)
+        # gather distances for token pairs using mapped nodes
+        idx_i = token_nodes.view(1, 1, T, 1).expand(B, H, T, 1)
+        idx_j = token_nodes.view(1, 1, 1, T).expand(B, H, 1, T)
+        d_geo_tok = d_geo_sq.unsqueeze(1).gather(2, idx_i.repeat(1,1,1,J)).gather(3, idx_j.repeat(1,1,J,1))
+        d_geo_tok = d_geo_tok.squeeze(-1).squeeze(-2)  # (B,H,T,T) approx
+        # lightcone cost: temporal + directional
+        if positions is None:
+            positions = torch.arange(T, device=Q.device).view(1, T).expand(B, T)
+        pos_i = positions.unsqueeze(1)  # (B,1,T)
+        pos_j = positions.unsqueeze(2)  # (B,T,1)
+        d_time = (pos_j - pos_i).abs().float() / max(1, self.params.tau)  # (B,T,T)
+        d_time = d_time.unsqueeze(1).expand(B, H, T, T)
+        if u_dir is None:
+            u_dir = torch.zeros(B, Dh, device=Q.device)
+        # direction penalty via 1 - cos(angle(u, ΔK)) as proxy
+        # ΔK ~ K_j (we ignore i to keep cheap): compute sim between u and K
+        Ku = F.normalize(K.mean(dim=2), dim=-1)  # (B,H,Dh)
+        u_n = F.normalize(u_dir, dim=-1).unsqueeze(1)  # (B,1,Dh)
+        dir_cost = (1.0 - (Ku * u_n).sum(-1, keepdim=True)).clamp_min(0.0)  # (B,H,1)
+        dir_cost = dir_cost.expand(B, H, T)  # (B,H,T)
+        dir_cost = dir_cost.unsqueeze(-1).expand(B, H, T, T)
+        # combine
+        scores = dots - self.params.beta.abs() * d_geo_tok - self.params.gamma.abs() * (
+            self.params.lambda_time.abs() * d_time + self.params.lambda_dir.abs() * dir_cost
+        )
+        return scores
+
+
+# ===============================
+# NTI Controller (episodic intent offset)
+# ===============================
+class NTIController(nn.Module):
+    def __init__(self, d: int, vocab_size: int, r_gamma: int = 32, offset_scale: float = 0.5, tau: int = 64):
+        super().__init__()
+        self.tau = tau
+        self.offset_scale = offset_scale
+        self.proj = nn.Sequential(
+            nn.Linear(d + r_gamma, 128), nn.GELU(),
+            nn.Linear(128, vocab_size)
+        )
+
+    def forward(self, H_seg: torch.Tensor, Gamma_seg: torch.Tensor,
+                attn_entropy: Optional[torch.Tensor] = None,
+                path_dev: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Compute Δz episodic offset for a segment.
+        H_seg: (B, τ, d), Gamma_seg: (B, τ, rγ)
+        attn_entropy/path_dev: optional scalars per batch
+        Return Δz: (B, V)
+        """
+        h_bar = H_seg.mean(dim=1)           # (B,d)
+        g_bar = Gamma_seg.mean(dim=1)       # (B,rγ)
+        x = torch.cat([h_bar, g_bar], dim=-1)
+        dz = self.proj(x) * self.offset_scale
+        return dz
+
+
+# ===============================
+# Top-level: N‑Transformers Coupler
+# ===============================
+@dataclass
+class NTCfg:
+    d: int = 2048
+    vocab_size: int = 50000
+    r_gamma: int = 32
+    J: int = 256
+    k: int = 16
+    K: int = 16
+    alpha: float = 0.05
+    noise_eps: float = 1e-3
+    kappa_thresh: float = 0.6
+    nti_tau: int = 64
+    nti_period: int = 16
+    offset_scale: float = 0.5
+
+
+class NTransformerCoupler(nn.Module):
+    """Parallel PF path + couplings to augment any decoder‑only LM.
+
+    Exposes step() for single‑step inference/training and segment_update() for NTI updates.
+    """
+    def __init__(self, cfg: NTCfg):
+        super().__init__()
+        self.cfg = cfg
+        pf_cfg = PFConfig(J=cfg.J, k=cfg.k, K=cfg.K, alpha=cfg.alpha, noise_eps=cfg.noise_eps)
+        self.adapter_out = PFAdapterOut(d=cfg.d, J=cfg.J, k=cfg.k)
+        self.ime = PFIntrinsicMetricEngine(k=cfg.k, hidden=128, metric_eps=1e-4)
+        self.pf_core = PFFieldCore(cfg=pf_cfg)
+        self.introspect = PFIntrospection(d=cfg.d, k=cfg.k, r_gamma=cfg.r_gamma)
+        self.integrator = PFIntegrationMeter(J=cfg.J, kappa_thresh=cfg.kappa_thresh)
+        self.gate = LogitGate(vocab_size=cfg.vocab_size, r_gamma=cfg.r_gamma)
+        self.lca = LCAWrapper(LCAParams(beta=0.7, gamma=0.3, lambda_time=0.2, lambda_dir=0.3, tau=cfg.nti_tau))
+        self.nti = NTIController(d=cfg.d, vocab_size=cfg.vocab_size, r_gamma=cfg.r_gamma,
+                                 offset_scale=cfg.offset_scale, tau=cfg.nti_tau)
+
+    def initial_state(self, batch_size: int, device: Optional[torch.device] = None) -> Dict[str, torch.Tensor]:
+        device = device or next(self.parameters()).device
+        F0 = torch.randn(batch_size, self.cfg.J, self.cfg.k, device=device) * 0.02
+        # Build initial KNN on PF nodes (use features themselves for init)
+        idx = knn_indices(F0, self.cfg.K)
+        A0 = build_adjacency(idx, self.cfg.J)
+        g0 = self.ime(F0)
+        return {"F": F0, "A": A0, "g": g0}
+
+    @torch.no_grad()
+    def rebuild_graph(self, F_t: torch.Tensor) -> torch.Tensor:
+        idx = knn_indices(F_t, self.cfg.K)
+        A = build_adjacency(idx, self.cfg.J)
+        return A
+
+    def step(self, state: Dict[str, torch.Tensor], h_t: torch.Tensor,
+             z_base_t: torch.Tensor,
+             Q: Optional[torch.Tensor] = None, K: Optional[torch.Tensor] = None,
+             positions: Optional[torch.Tensor] = None,
+             u_dir: Optional[torch.Tensor] = None) -> Tuple[Dict[str, torch.Tensor], torch.Tensor, Dict[str, torch.Tensor]]:
+        """Single decoding step coupling.
+        state: dict with F (B,J,k), A (B,J,J), g (B,k,k)
+        h_t: (B,d) token hidden; z_base_t: (B,V)
+        Q,K: attention tensors (B,H,T,dh) for optional LCA modulation; if None, gating only
+        Returns: new_state, z_final, logs
+        """
+        F_t, A_t, g_t = state["F"], state["A"], state["g"]
+        F_tilde = self.adapter_out(h_t)               # (B,J,k)
+        # evolve PF one step
+        F_next, pf_stats = self.pf_core(F_t, h_t, g_t, A_t, F_tilde)
+        # update metric and (optionally) graph
+        g_next = self.ime(F_next)
+        with torch.no_grad():
+            A_next = self.rebuild_graph(F_next)
+        # introspection & integration
+        # compute Syn/Conn
+        syn = self.integrator.synchrony(F_next, A_next)
+        conn = self.integrator.connectivity(A_next)
+        V, a, Gamma = self.introspect(F_next, h_t, syn, conn)
+        syn, conn, kappa, broadcast = self.integrator(F_next, A_next, V, a)
+        # gating logits
+        z_final = self.gate(z_base_t, Gamma)
+        # optional: LCA on attention scores (external model must accept these scores)
+        lca_scores = None
+        if Q is not None and K is not None:
+            lca_scores = self.lca(Q, K, F_next, g_next, positions=positions, u_dir=u_dir)
+        new_state = {"F": F_next, "A": A_next, "g": g_next,
+                     "V": V.detach(), "a": a.detach(), "Gamma": Gamma.detach(),
+                     "kappa": kappa.detach(), "broadcast": broadcast.detach()}
+        logs = {"tv": pf_stats["tv"], "omega": pf_stats["omega"],
+                "syn": syn.mean().detach(), "conn": conn.mean().detach(),
+                "V": V.mean().detach(), "a": a.mean().detach(), "kappa": kappa.mean().detach()}
+        if lca_scores is not None:
+            logs["lca_min"] = lca_scores.min().detach()
+            logs["lca_max"] = lca_scores.max().detach()
+        return new_state, z_final, logs
+
+    def segment_update(self, H_seg: torch.Tensor, Gamma_seg: torch.Tensor,
+                        attn_entropy: Optional[torch.Tensor] = None,
+                        path_dev: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Every r steps, compute episodic Δz via NTI.
+        H_seg: (B, τ, d); Gamma_seg: (B, τ, rγ)
+        Return: Δz (B,V) to be added to subsequent logits (late‑fusion)
+        """
+        dz = self.nti(H_seg, Gamma_seg, attn_entropy, path_dev)
+        return dz
+
+
+# ===============================
+# Losses (PF‑side; to be combined with LLM next‑token loss)
+# ===============================
+class PFLosses(nn.Module):
+    def __init__(self, lambda_coh: float = 0.5, lambda_gauge: float = 0.5,
+                 lambda_val: float = 0.2, lambda_self: float = 0.2, lambda_meta: float = 0.4):
+        super().__init__()
+        self.lambda_coh = lambda_coh
+        self.lambda_gauge = lambda_gauge
+        self.lambda_val = lambda_val
+        self.lambda_self = lambda_self
+        self.lambda_meta = lambda_meta
+
+    @staticmethod
+    def tv_loss(F_t: torch.Tensor, A: torch.Tensor) -> torch.Tensor:
+        Fi = F_t.unsqueeze(2)
+        Fj = F_t.unsqueeze(1)
+        diff = (Fi - Fj).pow(2).sum(-1).sqrt()  # (B,J,J)
+        return (diff * A).mean()
+
+    @staticmethod
+    def incoh_loss(H_t: torch.Tensor, F_t: torch.Tensor) -> torch.Tensor:
+        # penalize low alignment between pooled F and h
+        F_pool = F_t.mean(dim=1)
+        cos = F.cosine_similarity(F.normalize(F_pool, dim=-1), F.normalize(H_t, dim=-1))
+        return (1.0 - cos).mean()
+
+    @staticmethod
+    def pathdev_loss() -> torch.Tensor:
+        # placeholder; requires tracking best path proxy; return small constant to avoid zero grads
+        return torch.tensor(0.0, device=next(PFLosses().parameters()).device)
+
+    def forward(self, H_t: torch.Tensor, F_t: torch.Tensor, A_t: torch.Tensor,
+                V_t: torch.Tensor, a_t: torch.Tensor,
+                V_target: Optional[torch.Tensor] = None,
+                a_target: Optional[torch.Tensor] = None,
+                meta_pos: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+                meta_neg: Optional[Tuple[torch.Tensor, torch.Tensor]] = None) -> torch.Tensor:
+        # coherence
+        L_coh = self.tv_loss(F_t, A_t)
+        # gauge
+        L_gauge = self.incoh_loss(H_t, F_t) + self.pathdev_loss()
+        # valence (regression to target if provided)
+        if V_target is not None:
+            L_val = F.mse_loss(V_t, V_target)
+        else:
+            L_val = torch.zeros((), device=F_t.device)
+        # self/now
+        if a_target is not None:
+            L_self = F.binary_cross_entropy(a_t.clamp(1e-4, 1-1e-4), a_target)
+        else:
+            L_self = torch.zeros((), device=F_t.device)
+        # meta (contrastive on PF signatures; here use pooled F as proxy)
+        L_meta = torch.zeros((), device=F_t.device)
+        if (meta_pos is not None) and (meta_neg is not None):
+            F_pos1, F_pos2 = meta_pos  # both (B,J,k)
+            F_neg1, F_neg2 = meta_neg
+            # cosine distance between pooled features
+            def pool(Fx):
+                return F.normalize(Fx.mean(dim=1), dim=-1)
+            pos = 1.0 - F.cosine_similarity(pool(F_pos1), pool(F_pos2)).mean()
+            neg = F.cosine_similarity(pool(F_neg1), pool(F_neg2)).mean()
+            margin = 0.3
+            L_meta = F.relu(pos + neg - margin)
+        L = (self.lambda_coh * L_coh + self.lambda_gauge * L_gauge +
+             self.lambda_val * L_val + self.lambda_self * L_self + self.lambda_meta * L_meta)
+        return L
+
+
+# ===============================
+# Example Integration Skeleton (with a generic LM)
+# ===============================
+class DummyDecoderOnlyLM(nn.Module):
+    """Placeholder LM exposing hidden and logits for demonstration only.
+    Replace with your actual Transformer decoder (e.g., GPT‑like) and wire the coupler around it.
+    """
+    def __init__(self, d: int, vocab_size: int):
+        super().__init__()
+        self.d = d
+        self.emb = nn.Embedding(vocab_size, d)
+        self.ff = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, d))
+        self.head = nn.Linear(d, vocab_size)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        H = self.emb(x)  # (B,T,d)
+        H = self.ff(H)
+        logits = self.head(H)  # (B,T,V)
+        return H, logits
+
+
+class NTransformersModel(nn.Module):
+    """Full model wrapper: LM + N‑Transformers coupler.
+    This is a minimal training‑ready scaffold; extend as needed.
+    """
+    def __init__(self, lm: nn.Module, coupler: NTransformerCoupler, losses: PFLosses):
+        super().__init__()
+        self.lm = lm
+        self.coupler = coupler
+        self.losses = losses
+
+    def forward(self, x: torch.Tensor, y: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
+        B, T = x.shape
+        device = x.device
+        # initialize PF state
+        state = self.coupler.initial_state(B, device=device)
+        H, logits_base = self.lm(x)  # (B,T,d), (B,T,V)
+        logits = torch.empty_like(logits_base)
+        Gamma_hist = []
+        # step through time
+        for t in range(T):
+            h_t = H[:, t, :]
+            z_base_t = logits_base[:, t, :]
+            state, z_t, logs = self.coupler.step(state, h_t, z_base_t)
+            logits[:, t, :] = z_t
+            if "Gamma" in state:
+                Gamma_hist.append(state["Gamma"])  # (B,rγ)
+        # NTI every nti_period (here apply once at end for demo)
+        if len(Gamma_hist) >= self.coupler.cfg.nti_tau:
+            Gamma_seg = torch.stack(Gamma_hist[-self.coupler.cfg.nti_tau:], dim=1)  # (B,τ,rγ)
+            H_seg = H[:, -self.coupler.cfg.nti_tau:, :]
+            dz = self.coupler.segment_update(H_seg, Gamma_seg)
+            logits[:, -1, :] = logits[:, -1, :] + dz  # late‑fusion on final step (demo)
+        out = {"logits": logits}
+        if y is not None:
+            # next-token CE loss
+            loss_llm = F.cross_entropy(logits[:, :-1, :].reshape(-1, logits.size(-1)), y[:, 1:].reshape(-1))
+            # PF‑side loss (using last step as example)
+            H_last = H[:, -1, :]
+            F_last, A_last = state["F"], state["A"]
+            V_last, a_last = state["V"], state["a"]
+            loss_pf = self.losses(H_last, F_last, A_last, V_last, a_last)
+            loss = loss_llm + loss_pf
+            out.update({"loss": loss, "loss_llm": loss_llm, "loss_pf": loss_pf})
+        return out
+
+
+# ===============================
+# Quick smoke test (CPU)
+# ===============================
+if __name__ == "__main__":
+    torch.manual_seed(42)
+    cfg = NTCfg(d=256, vocab_size=8192, J=64, k=8, K=8, nti_tau=16, nti_period=8)
+    lm = DummyDecoderOnlyLM(d=cfg.d, vocab_size=cfg.vocab_size)
+    coupler = NTransformerCoupler(cfg)
+    losses = PFLosses()
+    model = NTransformersModel(lm, coupler, losses)
+
+    B, T = 4, 24
+    x = torch.randint(0, cfg.vocab_size, (B, T))
+    y = x.clone()
+
+    out = model(x, y)
+    print({k: float(v) if torch.is_tensor(v) and v.dim()==0 else v.shape for k, v in out.items() if k.startswith('loss') or k=='logits'})