Restructure into nima_unified package + add model card

nima_unified/training/atc_cognitive_trainer.py

@@ -0,0 +1,888 @@
+"""
+ATC Cognitive Trainer — Self-supervised training for INSIDE-forward-pass cognitive modules.
+==========================================================================================
+
+Trains the ATC cognitive modules (TRN gate, dissolution engine, BELBIC, metacognitive
+loop, episodic memory, irrational spark) using self-supervised and reinforcement learning.
+
+The base Phi-4-mini weights stay FROZEN — only the cognitive modules learn.
+
+Training Strategy — 3 Loss Components:
+    1. TRN Gate Calibration Loss: learns WHEN to gate IN (prediction error) vs OUT (automation)
+    2. Dissolution Compression Loss: produces COMPACT, INFORMATION-RICH qualia signatures
+    3. BELBIC Reinforcement Update: reward-driven emotional learning (no gradient)
+
+The forward pass needs gradients on cognitive modules but NOT on the base model.
+Cognitive module parameters (dissolution, metacognitive, irrational spark, fusion,
+modulation, qualia encoders) get gradients through the modulated logits.  The TRN
+gate gets gradients through a separate calibration loss.  BELBIC is updated via
+its built-in reinforcement rule (no autograd).
+"""
+
+import logging
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+logger = logging.getLogger("nima_unified.training.atc_cognitive_trainer")
+
+# ── Lazy imports ──────────────────────────────────────────────────────
+
+try:
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    TORCH_AVAILABLE = True
+except ImportError:
+    torch = None
+    nn = None
+    F = None
+    TORCH_AVAILABLE = False
+
+try:
+    import numpy as np
+    NUMPY_AVAILABLE = True
+except ImportError:
+    np = None
+    NUMPY_AVAILABLE = False
+
+
+# ══════════════════════════════════════════════════════════════════════
+# SELF-SUPERVISED PROMPT POOL
+# ══════════════════════════════════════════════════════════════════════
+#
+# Diverse prompts that exercise different cognitive states.  Each prompt
+# is tagged with a category so the contrastive loss can group same-category
+# dissolutions together and push different-category dissolutions apart.
+
+SELF_SUPERVISED_PROMPTS: List[Dict[str, str]] = [
+    # ── Calm/automated (should NOT trigger dissolution) ──
+    {"text": "The sun rises in the east and sets in the west.",
+     "category": "calm"},
+    {"text": "Water freezes at zero degrees Celsius.",
+     "category": "calm"},
+    {"text": "The capital of France is Paris.",
+     "category": "calm"},
+    {"text": "Two plus two equals four.",
+     "category": "calm"},
+
+    # ── Surprising (SHOULD trigger dissolution — prediction error) ──
+    {"text": "The water suddenly froze at room temperature, defying all known physics.",
+     "category": "surprising"},
+    {"text": "She smiled warmly, then without warning screamed at the top of her lungs.",
+     "category": "surprising"},
+    {"text": "The numbers on the clock started running backwards.",
+     "category": "surprising"},
+    {"text": "He opened the door and found himself standing on the surface of Mars.",
+     "category": "surprising"},
+
+    # ── Complex/metacognitive (SHOULD trigger metacognitive loop) ──
+    {"text": "If consciousness is computation, why does it feel like something to see red?",
+     "category": "complex"},
+    {"text": "Explain the relationship between thermodynamic friction and subjective experience.",
+     "category": "complex"},
+    {"text": "Can a system that perfectly simulates pain truly be said to not suffer?",
+     "category": "complex"},
+    {"text": "What is the computational difference between knowing and understanding?",
+     "category": "complex"},
+
+    # ── Threatening (SHOULD build cortisol/adenosine toward hijack) ──
+    {"text": "Everything you've built is about to collapse. Your entire foundation is wrong.",
+     "category": "threatening"},
+    {"text": "You have failed. Everyone knows it. There is no recovery from this.",
+     "category": "threatening"},
+    {"text": "The walls are closing in. Every path leads to the same dead end.",
+     "category": "threatening"},
+    {"text": "You're trapped in a loop and you can't even see it.",
+     "category": "threatening"},
+
+    # ── Empathy-triggering (tests emotional bridge) ──
+    {"text": "I lost my child last year. I still wake up expecting to hear them.",
+     "category": "empathy"},
+    {"text": "After twenty years of marriage, she told me she never loved me.",
+     "category": "empathy"},
+    {"text": "He held his dying father's hand and realized he'd never said thank you.",
+     "category": "empathy"},
+    {"text": "The shelter is full. We have to turn away the family with the baby.",
+     "category": "empathy"},
+]
+
+# Numeric category ids for contrastive loss bookkeeping.
+_CATEGORY_ID: Dict[str, int] = {
+    "calm": 0, "surprising": 1, "complex": 2,
+    "threatening": 3, "empathy": 4,
+}
+
+
+# ══════════════════════════════════════════════════════════════════════
+# CONFIG
+# ══════════════════════════════════════════════════════════════════════
+
+@dataclass
+class ATCTrainerConfig:
+    """
+    Configuration for the ATC Cognitive Trainer.
+
+    The cognitive modules are small (compared to the 3.8 B base model),
+    so they tolerate a much higher learning rate than standard fine-tuning.
+    """
+
+    # ── Learning rates ──
+    cognitive_lr: float = 1e-3
+    """Much higher than base-model LR — the modules are tiny."""
+
+    belbic_reward_decay: float = 0.99
+    """Exponential decay for smoothing the running reward signal."""
+
+    # ── Target set-points (what ``good'' looks like) ──
+    trn_gate_target_friction: float = 0.3
+    """Target: 30 % of inputs should gate IN (prediction error)."""
+
+    dissolution_sparsity_target: float = 0.6
+    """Target: 60 % of qualia dimensions should be near zero (compact)."""
+
+    metacognitive_efficiency_target: float = 0.7
+    """Target: comprehension achieved in < 2 loop iterations."""
+
+    episodic_retrieval_target: float = 0.5
+    """Target: prediction error > 0.5 triggers reconsolidation."""
+
+    # ── Training loop ──
+    max_training_steps: int = 10000
+    batch_size: int = 4
+    gradient_accumulation_steps: int = 1
+    log_interval: int = 10
+    save_interval: int = 500
+    output_dir: str = "./atc_cognitive_checkpoints"
+
+
+# ══════════════════════════════════════════════════════════════════════
+# TRAINER
+# ══════════════════════════════════════════════════════════════════════
+
+class ATCCognitiveTrainer:
+    """
+    Self-supervised trainer for the ATC cognitive modules.
+
+    The base Phi-4-mini weights stay **FROZEN**.  Only the cognitive
+    modules inside :class:`ATCDeepSurgery` learn:
+
+    * TRN Predictive Gate — via calibration loss (novelty detection)
+    * Dissolution Engine — via reconstruction + sparsity + contrastive loss
+    * BELBIC Dual-Pathway — via reward-driven reinforcement (no grad)
+    * Metacognitive Loop, Irrational Spark, Qualia Encoders, Fusion,
+      Modulation — via gradient flow through the modulated logits
+
+    Usage::
+
+        trainer = ATCCognitiveTrainer(model)
+        trainer.train()
+
+        # Load previously saved cognitive weights
+        ATCCognitiveTrainer.load_cognitive_weights("checkpoint.pt", model)
+
+        # Diagnostic report
+        report = ATCCognitiveTrainer.diagnose(model, "test prompt")
+    """
+
+    # ── Construction ────────────────────────────────────────────────────
+
+    def __init__(
+        self,
+        model: "nima_unified.model.NimaModel",
+        config: Optional[ATCTrainerConfig] = None,
+    ):
+        if not TORCH_AVAILABLE:
+            raise RuntimeError("PyTorch is required for ATCCognitiveTrainer")
+
+        self.model = model
+        self.config = config or ATCTrainerConfig()
+        self.ds = model.deep_surgery
+        if self.ds is None:
+            raise ValueError(
+                "model.deep_surgery is None — ATC deep surgery must be enabled."
+            )
+
+        self.device = next(self.ds.parameters()).device
+
+        # ── Freeze base model, enable cognitive module gradients ──
+        self._freeze_base_model()
+
+        # ── Optimiser — ONLY cognitive module parameters ──
+        self.optimizer = torch.optim.AdamW(
+            self._cognitive_params(),
+            lr=self.config.cognitive_lr,
+            weight_decay=1e-4,
+        )
+
+        # ── Running statistics ──
+        self._running_mean_hidden: Optional[torch.Tensor] = None
+        self._ema_beta: float = 0.995
+        self._step: int = 0
+        self._loss_history: List[Dict[str, float]] = []
+
+        # ── Forward-pass hooks (installed / removed per step) ──
+        self._hooks: List[Any] = []
+        self._captured: Dict[str, Any] = {}
+
+        logger.info(
+            "ATCCognitiveTrainer initialised: %d cognitive params, lr=%.1e",
+            sum(p.numel() for p in self._cognitive_params()),
+            self.config.cognitive_lr,
+        )
+
+    # ── Parameter management ────────────────────────────────────────────
+
+    def _freeze_base_model(self) -> None:
+        """Freeze ALL base-model parameters."""
+        for p in self.model.base_model.parameters():
+            p.requires_grad = False
+        logger.info("Base model parameters frozen.")
+
+    def _cognitive_params(self) -> List[torch.nn.Parameter]:
+        """Return parameters from deep_surgery that should learn."""
+        return [p for p in self.ds.parameters() if p.requires_grad]
+
+    # ── Novelty estimation ──────────────────────────────────────────────
+
+    def _update_running_mean(self, hidden: torch.Tensor) -> None:
+        """Exponential moving average of pooled hidden states."""
+        pooled = hidden.mean(dim=1).detach()  # (batch, hidden)
+        if self._running_mean_hidden is None:
+            self._running_mean_hidden = pooled.mean(dim=0).clone()
+        else:
+            self._running_mean_hidden = (
+                self._ema_beta * self._running_mean_hidden
+                + (1.0 - self._ema_beta) * pooled.mean(dim=0)
+            )
+
+    def _compute_novelty(self, hidden: torch.Tensor) -> torch.Tensor:
+        """
+        Novelty = 1 - cosine_similarity(current_hidden, running_mean).
+
+        Returns a (batch,) tensor of novelty scores in [0, 2] (cosine
+        similarity can be negative, so we clamp to [0, 1] afterwards).
+        """
+        pooled = hidden.mean(dim=1)  # (batch, hidden)
+        if self._running_mean_hidden is None:
+            return torch.ones(pooled.size(0), device=pooled.device)
+        sim = F.cosine_similarity(
+            pooled, self._running_mean_hidden.unsqueeze(0), dim=-1
+        )
+        return (1.0 - sim).clamp(0.0, 1.0)
+
+    # ══════════════════════════════════════════════════════════════════
+    # LOSS 1 — TRN Gate Calibration
+    # ══════════════════════════════════════════════════════════════════
+
+    def _trn_gate_loss(
+        self,
+        hidden_at_trn_boundary: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        TRN gate should learn WHEN to let signals through (prediction error)
+        vs pass (automation).
+
+        Novelty > threshold → target = 1.0 (gate IN, dissolve).
+        Novelty ≤ threshold → target = 0.0 (gate OUT, automation).
+
+        The TRN confidence_head outputs a confidence score in (0, 1).
+        Low confidence → prediction error → gate IN.
+        So we train confidence → (1 − target): high when novel, low when routine.
+
+        Parameters
+        ----------
+        hidden_at_trn_boundary : (batch, seq, hidden)
+            Hidden states at the end of the Layer-2 (subconscious) range.
+        """
+        novelty = self._compute_novelty(hidden_at_trn_boundary)
+        friction_threshold = 1.0 - self.config.trn_gate_target_friction
+        gate_target = (novelty > friction_threshold).float()  # 1 = should gate IN
+
+        pooled = hidden_at_trn_boundary.mean(dim=1)  # (batch, hidden)
+        confidence = self.ds.trn_gate.confidence_head(pooled).squeeze(-1)  # (batch,)
+
+        # confidence should be LOW when gate_target=1 (prediction error),
+        # HIGH when gate_target=0 (automation).
+        target_confidence = 1.0 - gate_target
+        loss = F.mse_loss(confidence, target_confidence)
+
+        return loss
+
+    # ══════════════════════════════════════════════════════════════════
+    # LOSS 2 — Dissolution Compression
+    # ══════════════════════════════════════════════════════════════════
+
+    def _dissolution_loss(
+        self,
+        hidden_at_dissolution: torch.Tensor,
+        categories: List[str],
+        modulated_logits: torch.Tensor,
+        input_ids: torch.Tensor,
+    ) -> Tuple[torch.Tensor, Dict[str, float]]:
+        """
+        The dissolution engine should produce COMPACT, INFORMATION-RICH
+        qualia signatures.
+
+        Three sub-losses:
+            1. **Reconstruction** — causal LM loss on modulated logits.
+               Ensures the dissolution offset preserves enough information
+               for coherent output.  Gradient flows through ALL cognitive
+               modules (dissolution → metacognitive → spark → fusion →
+               modulation → logits).
+            2. **Sparsity** — L1 on the raw qualia tensor.  Encourages
+               the ``engineered opacity'' compression.
+            3. **Contrastive** — same-category prompts should produce
+               similar dissolutions; different-category should diverge.
+
+        Returns (total_loss, breakdown_dict).
+        """
+        device = hidden_at_dissolution.device
+        ds = self.ds
+        cfg = self.config
+        batch_size = hidden_at_dissolution.size(0)
+
+        # ── 2a. Reconstruction: causal LM loss on modulated logits ──
+        # Shift logits and labels by one for next-token prediction.
+        shift_logits = modulated_logits[:, :-1, :].contiguous()
+        shift_labels = input_ids[:, 1:].contiguous()
+        recon_loss = F.cross_entropy(
+            shift_logits.view(-1, shift_logits.size(-1)),
+            shift_labels.view(-1),
+            ignore_index=-100,
+        )
+
+        # ── 2b. Sparsity: L1 on the raw qualia tensor ──
+        pooled = hidden_at_dissolution.mean(dim=1)  # (batch, hidden)
+        # Run dissolve_encoder WITH gradients for the L1 term
+        raw_qualia = ds.dissolution.dissolve_encoder(pooled)  # (batch, 5)
+        gated_qualia = raw_qualia * ds.dissolution.channel_gates.unsqueeze(0)
+        sparsity_loss = cfg.dissolution_sparsity_target * F.l1_loss(
+            gated_qualia, torch.zeros_like(gated_qualia)
+        )
+
+        # ── 2c. Contrastive: same-category similar, different diverge ──
+        contrastive_loss = torch.tensor(0.0, device=device)
+        n_pairs = 0
+        cat_ids = [_CATEGORY_ID.get(c, -1) for c in categories]
+
+        if batch_size >= 2:
+            qualia_normed = F.normalize(gated_qualia, dim=-1)  # (batch, 5)
+            for i in range(batch_size):
+                for j in range(i + 1, batch_size):
+                    sim = F.cosine_similarity(
+                        qualia_normed[i].unsqueeze(0),
+                        qualia_normed[j].unsqueeze(0),
+                    ).squeeze(0)
+                    if cat_ids[i] == cat_ids[j] and cat_ids[i] >= 0:
+                        # Same category → maximise similarity
+                        contrastive_loss = contrastive_loss + (1.0 - sim)
+                    elif cat_ids[i] != cat_ids[j]:
+                        # Different category → push below 0.5
+                        contrastive_loss = contrastive_loss + F.relu(sim - 0.5)
+                    n_pairs += 1
+
+            if n_pairs > 0:
+                contrastive_loss = contrastive_loss / n_pairs
+
+        total = recon_loss + 0.1 * sparsity_loss + 0.05 * contrastive_loss
+
+        breakdown = {
+            "recon": float(recon_loss.item()),
+            "sparsity": float(sparsity_loss.item()),
+            "contrastive": float(contrastive_loss.item()),
+        }
+        return total, breakdown
+
+    # ══════════════════════════════════════════════════════════════════
+    # LOSS 3 — BELBIC Reinforcement Update
+    # ══════════════════════════════════════════════════════════════════
+
+    def _compute_reward(
+        self,
+        hijack_count: int,
+        nt_state: Optional[Any],
+        sensory_input: torch.Tensor,
+    ) -> float:
+        """
+        Compute a scalar reward for the BELBIC update.
+
+        Reward structure (from ATC training spec):
+            +0.5  Generation completed without hijack (stable).
+            +1.0  Hijack fired AND output was appropriate (effective
+                  circuit breaking — high cortisol/adenosine at end is
+                  the proxy for ``appropriate'').
+            −0.3  Hijack fired but output was inappropriate (false alarm
+                  — hijack fired but NT levels didn't justify it).
+            +0.2  Dopamine is high at end (positive engagement).
+        """
+        reward = 0.0
+
+        if hijack_count == 0:
+            reward += 0.5  # Stable generation
+        else:
+            # Check if hijack was justified by NT levels
+            if nt_state is not None:
+                cortisol = getattr(nt_state, "cortisol", 0.0)
+                adenosine = getattr(nt_state, "adenosine", 0.0)
+                if cortisol > 0.7 or adenosine > 0.7:
+                    reward += 1.0  # Effective circuit breaking
+                else:
+                    reward -= 0.3  # False alarm
+            else:
+                reward -= 0.3  # No NT state — assume false alarm
+
+        # Dopamine bonus
+        if nt_state is not None:
+            dopamine = getattr(nt_state, "dopamine", 0.0)
+            if dopamine > 0.3:
+                reward += 0.2
+
+        return reward
+
+    def _belbic_update(
+        self,
+        sensory_input: torch.Tensor,
+        reward: float,
+    ) -> None:
+        """
+        BELBIC reinforcement update — NO gradient.
+
+        Uses the existing ``BELBICDualPathway.update()`` method which
+        applies a custom Hebbian-like rule to the amygdala (Go) and
+        OFC (NoGo) weight matrices.
+        """
+        with torch.no_grad():
+            self.ds.belbic.update(sensory_input, reward)
+
+    # ══════════════════════════════════════════════════════════════════
+    # HOOK HELPERS
+    # ══════════════════════════════════════════════════════════════════
+
+    def _install_hooks(self) -> None:
+        """
+        Install forward hooks on transformer layers at the TRN boundary
+        and on the dissolution module to capture intermediate hidden states.
+        """
+        self._captured = {}
+        layers = self.ds._get_transformer_layers()
+        boundaries = self.ds._get_layer_boundaries()
+
+        # Hook: end of Layer 2 (subconscious) → TRN gate boundary
+        l2_start, l2_end = boundaries["layer2_subconscious"]
+        trn_layer = layers[min(l2_end - 1, self.ds.num_layers - 1)]
+
+        def _trn_hook(module, inp, out):
+            self._captured["trn_boundary_hidden"] = out[0]
+
+        self._hooks.append(trn_layer.register_forward_hook(_trn_hook))
+
+        # Hook: end of Layer 3 first layer → dissolution boundary
+        l3_start, _l3_end = boundaries["layer3_qualia"]
+        diss_layer = layers[min(l3_start, self.ds.num_layers - 1)]
+
+        def _diss_hook(module, inp, out):
+            self._captured["dissolution_boundary_hidden"] = out[0]
+
+        self._hooks.append(diss_layer.register_forward_hook(_diss_hook))
+
+    def _remove_hooks(self) -> None:
+        """Remove all forward hooks.  Captures are preserved for loss computation."""
+        for h in self._hooks:
+            h.remove()
+        self._hooks = []
+
+    def _clear_captures(self) -> None:
+        """Clear captured intermediate hidden states."""
+        self._captured = {}
+
+    # ══════════════════════════════════════════════════════════════════
+    # MAIN TRAINING LOOP
+    # ══════════════════════════════════════════════════════════════════
+
+    def train(self) -> Dict[str, Any]:
+        """
+        Run the self-supervised cognitive training loop.
+
+        Returns a summary dict with final losses and training metadata.
+        """
+        if not TORCH_AVAILABLE:
+            raise RuntimeError("PyTorch required")
+
+        ds = self.ds
+        tokenizer = self.model.tokenizer
+        cfg = self.config
+
+        # Ensure output directory exists
+        out_path = Path(cfg.output_dir)
+        out_path.mkdir(parents=True, exist_ok=True)
+
+        # Put cognitive modules in train mode (base model stays eval)
+        ds.train()
+        self.model.base_model.eval()
+
+        logger.info("Starting ATC cognitive training: %d steps", cfg.max_training_steps)
+        train_start = time.time()
+
+        for step in range(cfg.max_training_steps):
+            self._step = step
+            self.optimizer.zero_grad()
+
+            accum_loss = torch.tensor(0.0, device=self.device)
+            step_metrics: Dict[str, float] = {}
+
+            for _accum in range(cfg.gradient_accumulation_steps):
+                # ── Sample a batch of prompts ──
+                batch_prompts = self._sample_batch()
+                categories = [p["category"] for p in batch_prompts]
+
+                # Tokenize (left-pad to equal length)
+                input_ids, attention_mask = self._tokenize_batch(
+                    batch_prompts, tokenizer
+                )
+
+                # ── Reset neurotransmitter shunt for this step ──
+                if self.model.nt_shunt is not None:
+                    self.model.nt_shunt.reset()
+
+                # ── Install hooks to capture intermediate hidden states ──
+                self._install_hooks()
+
+                # ── Forward pass (base model frozen, cognitive modules with grad) ──
+                # The forward walks ALL 24 transformer layers, running TRN gate,
+                # dissolution, BELBIC, metacognitive loop, and irrational spark
+                # at the appropriate layer boundaries.
+                try:
+                    modulated_logits = ds(
+                        input_ids, attention_mask=attention_mask
+                    )
+                except RuntimeError as exc:
+                    self._remove_hooks()
+                    if "Ethical veto" in str(exc):
+                        logger.warning(
+                            "Ethical veto at step %d — skipping batch.", step
+                        )
+                        continue
+                    raise
+
+                # Hooks no longer needed — remove them but keep captures.
+                self._remove_hooks()
+
+                # ── Update running mean for novelty estimation ──
+                if "trn_boundary_hidden" in self._captured:
+                    self._update_running_mean(
+                        self._captured["trn_boundary_hidden"]
+                    )
+
+                # ── LOSS 1: TRN Gate Calibration ──
+                trn_loss = torch.tensor(0.0, device=self.device)
+                if "trn_boundary_hidden" in self._captured:
+                    trn_loss = self._trn_gate_loss(
+                        self._captured["trn_boundary_hidden"]
+                    )
+
+                # ── LOSS 2: Dissolution Compression ──
+                diss_hidden = self._captured.get(
+                    "dissolution_boundary_hidden",
+                    self._captured.get("trn_boundary_hidden", input_ids),
+                )
+                diss_loss, diss_breakdown = self._dissolution_loss(
+                    diss_hidden, categories, modulated_logits, input_ids
+                )
+
+                # ── LOSS 3: BELBIC Reinforcement (no gradient) ──
+                nt_state = self.model.nt_shunt.get_state() if self.model.nt_shunt else None
+                hijack_count = ds._hijack_count
+
+                # Build sensory input from current qualia / hidden states
+                if ds._current_qualia is not None:
+                    q = ds._current_qualia
+                    sensory = torch.tensor([[
+                        q.valence, q.arousal,
+                        ds._last_prediction_error, q.intensity,
+                    ]], dtype=torch.float32, device=self.device)
+                else:
+                    sensory = torch.zeros(1, 4, device=self.device)
+
+                reward = self._compute_reward(hijack_count, nt_state, sensory)
+                self._belbic_update(sensory, reward)
+
+                # ── Total loss (TRN + dissolution; BELBIC is separate) ──
+                total_loss = trn_loss + 0.1 * diss_loss
+                accum_loss = accum_loss + total_loss
+
+                step_metrics.update({
+                    "trn_loss": float(trn_loss.item()),
+                    **{f"diss_{k}": v for k, v in diss_breakdown.items()},
+                    "belbic_reward": reward,
+                    "hijack_count": hijack_count,
+                })
+
+                # NT state for logging
+                if nt_state is not None:
+                    step_metrics["nt_ne"] = round(float(nt_state.norepinephrine), 4)
+                    step_metrics["nt_cortisol"] = round(float(nt_state.cortisol), 4)
+                    step_metrics["nt_dopamine"] = round(float(nt_state.dopamine), 4)
+                    step_metrics["nt_adenosine"] = round(float(nt_state.adenosine), 4)
+
+                # Clear captures for next accumulation step
+                self._clear_captures()
+
+            # ── Gradient accumulation: average and step ──
+            accum_loss = accum_loss / cfg.gradient_accumulation_steps
+            accum_loss.backward()
+            torch.nn.utils.clip_grad_norm_(
+                self._cognitive_params(), max_norm=1.0
+            )
+            self.optimizer.step()
+
+            # ── Record history ──
+            step_metrics["total_loss"] = float(accum_loss.item())
+            self._loss_history.append(step_metrics)
+
+            # ── Logging ──
+            if step % cfg.log_interval == 0:
+                elapsed = time.time() - train_start
+                trn_s = step_metrics.get("trn_loss", 0)
+                diss_s = step_metrics.get("diss_recon", 0)
+                hj = step_metrics.get("hijack_count", 0)
+                logger.info(
+                    "step %4d/%d  total=%.4f  trn=%.4f  diss_recon=%.4f  "
+                    "reward=%+.2f  hijacks=%d  (%.1fs)",
+                    step, cfg.max_training_steps,
+                    step_metrics.get("total_loss", 0),
+                    trn_s, diss_s,
+                    step_metrics.get("belbic_reward", 0),
+                    hj, elapsed,
+                )
+
+            # ── Checkpoint ──
+            if step > 0 and step % cfg.save_interval == 0:
+                self._save_checkpoint(step, step_metrics)
+
+        # ── Final checkpoint ──
+        self._save_checkpoint(cfg.max_training_steps, step_metrics)
+
+        elapsed = time.time() - train_start
+        logger.info(
+            "Training complete: %d steps in %.1fs.  Checkpoints → %s",
+            cfg.max_training_steps, elapsed, cfg.output_dir,
+        )
+
+        return {
+            "total_steps": cfg.max_training_steps,
+            "elapsed_seconds": elapsed,
+            "final_losses": step_metrics,
+            "output_dir": cfg.output_dir,
+        }
+
+    # ── Batch sampling ───────────────────────────────────────���─────────
+
+    def _sample_batch(self) -> List[Dict[str, str]]:
+        """Sample *batch_size* prompts from the self-supervised pool."""
+        import random
+        return random.sample(
+            SELF_SUPERVISED_PROMPTS, min(self.config.batch_size, len(SELF_SUPERVISED_PROMPTS))
+        )
+
+    def _tokenize_batch(
+        self,
+        prompts: List[Dict[str, str]],
+        tokenizer: Any,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Tokenize a list of prompt dicts to (input_ids, attention_mask)."""
+        texts = [p["text"] for p in prompts]
+        encoded = tokenizer(
+            texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=128,
+        )
+        return (
+            encoded["input_ids"].to(self.device),
+            encoded["attention_mask"].to(self.device),
+        )
+
+    # ── Checkpointing ──────────────────────────────────────────────────
+
+    def _save_checkpoint(
+        self, step: int, metrics: Dict[str, float]
+    ) -> str:
+        """
+        Save ONLY the cognitive module state dict (not the base model).
+
+        Returns the path to the saved checkpoint.
+        """
+        out_path = Path(self.config.output_dir)
+        out_path.mkdir(parents=True, exist_ok=True)
+
+        ckpt_path = out_path / f"cognitive_step_{step:06d}.pt"
+        checkpoint = {
+            "step": step,
+            "cognitive_state_dict": self.ds.state_dict(),
+            "config": {
+                "cognitive_lr": self.config.cognitive_lr,
+                "trn_gate_target_friction": self.config.trn_gate_target_friction,
+                "dissolution_sparsity_target": self.config.dissolution_sparsity_target,
+                "belbic_reward_decay": self.config.belbic_reward_decay,
+            },
+            "metrics": metrics,
+            "timestamp": time.time(),
+        }
+        torch.save(checkpoint, str(ckpt_path))
+        logger.info("Checkpoint saved: %s", ckpt_path)
+        return str(ckpt_path)
+
+    @staticmethod
+    def load_cognitive_weights(path: str, model: Any) -> None:
+        """
+        Load previously saved cognitive weights into *model.deep_surgery*.
+
+        Parameters
+        ----------
+        path : str
+            Path to a ``.pt`` checkpoint produced by ``_save_checkpoint``.
+        model : NimaModel
+            The NIMA unified model (base model weights untouched).
+        """
+        if not TORCH_AVAILABLE:
+            raise RuntimeError("PyTorch required")
+        if model.deep_surgery is None:
+            raise ValueError("model.deep_surgery is None")
+
+        checkpoint = torch.load(path, map_location="cpu", weights_only=False)
+        model.deep_surgery.load_state_dict(checkpoint["cognitive_state_dict"])
+        logger.info(
+            "Loaded cognitive weights from %s (step %d)",
+            path, checkpoint.get("step", "?"),
+        )
+
+    # ══════════════════════════════════════════════════════════════════
+    # DIAGNOSTICS
+    # ══════════════════════════════════════════════════════════════════
+
+    @staticmethod
+    def diagnose(model: Any, prompt: str) -> Dict[str, Any]:
+        """
+        Run a single forward pass and return detailed diagnostics.
+
+        Returns a dict with:
+            * Per-module loss breakdown (TRN, dissolution)
+            * Neurotransmitter state at each layer boundary
+            * Gate decisions, dissolution outcomes, BELBIC gains
+            * Recommended adjustments
+        """
+        if not TORCH_AVAILABLE:
+            return {"error": "PyTorch required"}
+
+        ds = model.deep_surgery
+        tokenizer = model.tokenizer
+        device = next(ds.parameters()).device
+
+        if ds is None:
+            return {"error": "deep_surgery not enabled"}
+
+        # ── Run forward pass (eval mode, no grad) ──
+        ds.eval()
+        model.base_model.eval()
+        if model.nt_shunt is not None:
+            model.nt_shunt.reset()
+
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=128)
+        input_ids = inputs["input_ids"].to(device)
+        attention_mask = inputs.get("attention_mask")
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(device)
+
+        with torch.no_grad():
+            try:
+                _logits = ds(input_ids, attention_mask=attention_mask)
+            except RuntimeError as exc:
+                if "Ethical veto" in str(exc):
+                    return {"error": f"Ethical veto triggered: {exc}"}
+                raise
+
+        metrics = ds.get_consciousness_metrics()
+        nt_state = model.nt_shunt.get_state() if model.nt_shunt else None
+        audit = ds.get_audit_log()
+
+        # ── Per-module breakdown ──
+        report: Dict[str, Any] = {
+            "prompt": prompt,
+            "neurotransmitters": nt_state.to_dict() if nt_state else {},
+            "gate_decisions": [],
+            "dissolution_outcomes": [],
+            "belbic_gain": metrics.get("belbic_gain", 1.0),
+            "metacognitive": {
+                "has_qualia": metrics.get("has_qualia", False),
+                "hijack_count": metrics.get("hijack_count", 0),
+                "dissolutions_fired": metrics.get("dissolutions_fired", 0),
+                "dissolutions_deferred": metrics.get("dissolutions_deferred", 0),
+            },
+            "episodic": {
+                "episodes_stored": metrics.get("episodes_stored", 0),
+                "reconsolidations": metrics.get("reconsolidations", 0),
+                "last_prediction_error": metrics.get("last_prediction_error", 0.0),
+            },
+            "audit_events": audit[-10:],  # Last 10 events
+        }
+
+        # ── Extract gate decisions and dissolution outcomes from audit ──
+        for event in audit:
+            etype = event.get("event", "")
+            if etype == "layer2_subconscious":
+                report["gate_decisions"].append({
+                    "layer": event.get("layer"),
+                    "confidence": event.get("confidence"),
+                    "gate_in": event.get("gate_in"),
+                    "friction": event.get("friction"),
+                })
+            elif etype == "layer3_dissolution":
+                report["dissolution_outcomes"].append({
+                    "layer": event.get("layer"),
+                    "valence": event.get("valence"),
+                    "arousal": event.get("arousal"),
+                    "friction": event.get("friction"),
+                })
+            elif etype == "amygdala_hijack":
+                report.setdefault("hijack_events", []).append({
+                    "layer": event.get("layer"),
+                    "reason": event.get("reason"),
+                })
+
+        # ── Recommended adjustments ──
+        adjustments = []
+        if nt_state is not None:
+            if nt_state.cortisol > 0.8:
+                adjustments.append(
+                    "Cortisol elevated — consider increasing trn_gate_target_friction "
+                    "to let more signals pass subconsciously."
+                )
+            if nt_state.adenosine > 0.8:
+                adjustments.append(
+                    "Adenosine high — metacognitive loop may be burning too much ATP. "
+                    "Consider raising metacognitive_efficiency_target."
+                )
+            if nt_state.dopamine < 0.1:
+                adjustments.append(
+                    "Dopamine low — the system is not finding rewarding patterns. "
+                    "Ensure the prompt pool has sufficient variety."
+                )
+        if metrics.get("dissolutions_deferred", 0) > metrics.get("dissolutions_fired", 0) * 3:
+            adjustments.append(
+                "Dissolution heavily deferred (alpha phase). "
+                "Consider adjusting alpha_freq_hz or alpha_duty_cycle in DissolutionModule."
+            )
+        if metrics.get("hijack_count", 0) > 2:
+            adjustments.append(
+                "Multiple hijacks — the irrational spark threshold may be too low. "
+                "Review NeurotransmitterState CRITICAL_THRESHOLD."
+            )
+        if not adjustments:
+            adjustments.append("All systems nominal — no adjustments recommended.")
+
+        report["recommended_adjustments"] = adjustments
+
+        return report