burt-imma / train_ablation.py

push from SNAPKITTYWEST/burt-imma

b88c26d verified about 2 months ago

24 kB

#!/usr/bin/env python3
"""
BURT-IMMA Full Training Script with Ablation Support
License: BSL-1.1
Contact: jessica@collectivekitty.com

Loads config from config/ablation_arithmetic.yaml and runs MMEP training
with ablation variants.

Usage:
  python train_ablation.py --config config/ablation_arithmetic.yaml
  python train_ablation.py --ablation no_memory --device cuda:0
  python train_ablation.py --ablation all --output-dir results/
"""

import argparse
import json
import os
import time
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np
import yaml

try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False
    raise ImportError("PyTorch is required for training. Install with: pip install torch")


# ---------------------------------------------------------------------------
# Model Components
# ---------------------------------------------------------------------------


class SmoothLeakyActivation(nn.Module):
    """alpha * x + (1-alpha) * x * sigmoid(x)"""

    def __init__(self, alpha: float = 0.01):
        super().__init__()
        self.alpha = alpha

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.alpha * x + (1.0 - self.alpha) * x * torch.sigmoid(x)


class GatesNormalization(nn.Module):
    """Constrained softmax with entropy <= max_entropy."""

    def __init__(self, d_input: int, num_gates: int, max_entropy: float = 0.20):
        super().__init__()
        self.gate_proj = nn.Linear(d_input, num_gates)
        self.max_entropy = max_entropy

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        logits = self.gate_proj(x)
        probs = F.softmax(logits, dim=-1)
        entropy = -(probs * (probs + 1e-10).log()).sum(dim=-1)

        # Sharpen via temperature reduction if entropy exceeds bound
        violations = entropy > self.max_entropy
        if violations.any():
            tau_lo = torch.full_like(entropy, 0.01)
            tau_hi = torch.ones_like(entropy)
            for _ in range(20):
                tau_mid = (tau_lo + tau_hi) / 2.0
                p = F.softmax(logits / tau_mid.unsqueeze(-1), dim=-1)
                h = -(p * (p + 1e-10).log()).sum(dim=-1)
                too_high = h > self.max_entropy
                tau_hi = torch.where(too_high, tau_mid, tau_hi)
                tau_lo = torch.where(too_high, tau_lo, tau_mid)
            tau_final = (tau_lo + tau_hi) / 2.0
            probs_sharp = F.softmax(logits / tau_final.unsqueeze(-1), dim=-1)
            probs = torch.where(violations.unsqueeze(-1), probs_sharp, probs)
            entropy = -(probs * (probs + 1e-10).log()).sum(dim=-1)

        return probs, entropy


class CIFGMatrixMemory(nn.Module):
    """Coupled Input-Forget Gate Matrix Memory."""

    def __init__(self, d_mem: int):
        super().__init__()
        self.d_mem = d_mem
        self.W_f = nn.Linear(d_mem, 1)
        self.register_buffer("C_global", torch.zeros(d_mem, d_mem))

    def forward(self, x: torch.Tensor, C_prev: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        f = torch.sigmoid(self.W_f(x))  # (batch, 1)
        candidate = x.unsqueeze(-1) @ x.unsqueeze(-2)  # (batch, d, d)
        candidate = candidate / (candidate.norm(dim=(-2, -1), keepdim=True) + 1e-8)
        C_new = f.unsqueeze(-1) * C_prev + (1.0 - f).unsqueeze(-1) * candidate
        return C_new, f

    def read(self, query: torch.Tensor, C: torch.Tensor) -> torch.Tensor:
        return torch.bmm(C, query.unsqueeze(-1)).squeeze(-1)


class GatesRouter(nn.Module):
    """Mixture-of-Experts router with top-k gating."""

    def __init__(self, d_input: int, num_experts: int, top_k: int = 1):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gate = nn.Linear(d_input, num_experts)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Returns (expert_indices, expert_weights, load_balance_loss)."""
        scores = self.gate(x)
        probs = F.softmax(scores, dim=-1)

        top_k_vals, top_k_idx = torch.topk(probs, self.top_k, dim=-1)
        top_k_vals = top_k_vals / (top_k_vals.sum(dim=-1, keepdim=True) + 1e-8)

        # Load balancing auxiliary loss
        routing_probs = probs.mean(dim=0)
        uniform = torch.ones(self.num_experts, device=x.device) / self.num_experts
        lb_loss = ((routing_probs - uniform) ** 2).sum()

        return top_k_idx, top_k_vals, lb_loss


class SuperpositionedInductionHeads(nn.Module):
    """Multi-head attention with induction bias."""

    def __init__(self, d_model: int, n_heads: int = 4):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch, seq_len, d = x.shape
        Q = self.W_Q(x).view(batch, seq_len, self.n_heads, self.d_head).transpose(1, 2)
        K = self.W_K(x).view(batch, seq_len, self.n_heads, self.d_head).transpose(1, 2)
        V = self.W_V(x).view(batch, seq_len, self.n_heads, self.d_head).transpose(1, 2)

        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_head ** 0.5)

        # Causal mask
        mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device) * float("-inf"), diagonal=1)
        scores = scores + mask.unsqueeze(0).unsqueeze(0)

        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)
        out = out.transpose(1, 2).contiguous().view(batch, seq_len, d)
        return self.W_O(out)


class ExpertFFN(nn.Module):
    """Single expert feed-forward network."""

    def __init__(self, d_model: int, d_ff: Optional[int] = None):
        super().__init__()
        d_ff = d_ff or d_model * 4
        self.W1 = nn.Linear(d_model, d_ff)
        self.W2 = nn.Linear(d_ff, d_model)
        self.activation = SmoothLeakyActivation()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.W2(self.activation(self.W1(x)))


class BURTIMMALayer(nn.Module):
    """Single BURT-IMMA layer."""

    def __init__(self, d_model: int, num_experts: int, top_k: int, d_mem: int):
        super().__init__()
        self.attention = SuperpositionedInductionHeads(d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.router = GatesRouter(d_model, num_experts, top_k)
        self.experts = nn.ModuleList([ExpertFFN(d_model) for _ in range(num_experts)])
        self.memory = CIFGMatrixMemory(d_mem)

    def forward(self, x: torch.Tensor, C_prev: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        # Attention with residual
        h = self.norm1(x + self.attention(x))

        # MoE routing
        batch, seq_len, d = h.shape
        h_flat = h.view(batch * seq_len, d)
        indices, weights, lb_loss = self.router(h_flat)

        # Dispatch to experts
        expert_out = torch.zeros_like(h_flat)
        for k in range(indices.shape[1]):
            for e_id in range(len(self.experts)):
                mask = (indices[:, k] == e_id)
                if mask.any():
                    out_e = self.experts[e_id](h_flat[mask])
                    expert_out[mask] += weights[mask, k:k+1] * out_e

        expert_out = expert_out.view(batch, seq_len, d)
        h = self.norm2(h + expert_out)

        # Memory update (use mean across sequence for memory key)
        h_mean = h.mean(dim=1)  # (batch, d)
        C_new, f_gate = self.memory(h_mean, C_prev)

        return h, C_new, lb_loss


class BURTIMMA(nn.Module):
    """Full BURT-IMMA model."""

    def __init__(self, hidden_dim: int, num_layers: int, num_experts: int,
                 top_k: int, d_mem: int):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.d_mem = d_mem

        self.input_proj = nn.Linear(hidden_dim, hidden_dim)
        self.layers = nn.ModuleList([
            BURTIMMALayer(hidden_dim, num_experts, top_k, d_mem)
            for _ in range(num_layers)
        ])
        self.gates_norm = GatesNormalization(hidden_dim, num_experts)
        self.output_proj = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        batch = x.shape[0]
        h = self.input_proj(x)

        C = torch.zeros(batch, self.d_mem, self.d_mem, device=x.device)
        total_lb_loss = torch.tensor(0.0, device=x.device)

        for layer in self.layers:
            h, C, lb_loss = layer(h, C)
            total_lb_loss = total_lb_loss + lb_loss

        # Gates normalization on final hidden
        h_mean = h.mean(dim=1)
        gate_probs, entropy = self.gates_norm(h_mean)

        out = self.output_proj(h)

        aux = {
            "entropy": entropy,
            "gate_probs": gate_probs,
            "lb_loss": total_lb_loss,
            "memory_trace": C.diagonal(dim1=-2, dim2=-1).sum(dim=-1),
        }
        return out, aux


# ---------------------------------------------------------------------------
# Ablation Trainer
# ---------------------------------------------------------------------------


class AblationTrainer:
    """Full ablation training with MMEP support."""

    def __init__(self, config: dict, ablation_name: str = "full_mmep",
                 device: str = "cpu", output_dir: str = "results"):
        self.config = config
        self.ablation_name = ablation_name
        self.device = torch.device(device)
        self.output_dir = Path(output_dir) / ablation_name
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # Extract config sections
        self.model_cfg = config["model"]
        self.mmep_cfg = config["mmep"]
        self.train_cfg = config["training"]
        self.data_cfg = config["data"]

        # Apply ablation overrides
        self.use_mmep = True
        self.use_memory = True
        self.use_constraint = True
        self.use_moe = True
        self._apply_ablation(ablation_name)

        # Build model and optimizer
        self.model = self.build_model()
        self.model.to(self.device)
        self.optimizer = torch.optim.Adam(
            self.model.parameters(),
            lr=self.mmep_cfg["lr_W"]
        )

        # Metrics log
        self.metrics_history: List[Dict[str, float]] = []

    def _apply_ablation(self, name: str) -> None:
        """Apply ablation variant configuration."""
        for ablation in self.config.get("ablations", []):
            if ablation["name"] == name:
                disable = ablation.get("disable", [])
                if "mmep" in disable:
                    self.use_mmep = False
                if "C_global" in disable or "C_expert" in disable:
                    self.use_memory = False
                if "projection" in disable:
                    self.use_constraint = False
                if ablation.get("use_backprop", False):
                    self.use_mmep = False
                if "override" in ablation:
                    for k, v in ablation["override"].items():
                        self.model_cfg[k] = v
                break

    def build_model(self) -> BURTIMMA:
        """Create BURT-IMMA model from config."""
        num_experts = self.model_cfg["num_experts"] if self.use_moe else 1
        return BURTIMMA(
            hidden_dim=self.model_cfg["hidden_dim"],
            num_layers=self.model_cfg["num_layers"],
            num_experts=num_experts,
            top_k=self.model_cfg["top_k"],
            d_mem=self.model_cfg["d_mem"],
        )

    def mmep_free_phase(self, x: torch.Tensor) -> torch.Tensor:
        """Run T_free relaxation steps (free phase of MMEP)."""
        T_free = self.mmep_cfg["T_free"]
        alpha = self.mmep_cfg["alpha_relax"]

        h = x.clone()
        for t in range(T_free):
            out, aux = self.model(h)
            # Relaxation: blend output back toward input
            h = (1.0 - alpha) * h + alpha * out

        return h

    def mmep_nudge_phase(self, x: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """Run T_nudge nudged steps (nudged phase of MMEP)."""
        T_nudge = self.mmep_cfg["T_nudge"]
        beta = self.mmep_cfg["beta"]
        alpha = self.mmep_cfg["alpha_relax"]

        h = x.clone()
        for t in range(T_nudge):
            out, aux = self.model(h)
            h = (1.0 - alpha) * h + alpha * out
            # Nudge toward target
            h = h + beta * (target - h)

        return h

    def compute_local_update(self, h_free: torch.Tensor,
                             h_nudged: torch.Tensor) -> torch.Tensor:
        """Compute local parameter update: (nudged - free) / beta."""
        beta = self.mmep_cfg["beta"]
        return (h_nudged - h_free) / beta

    def _spectral_norm_project(self) -> None:
        """Project weight matrices to satisfy spectral norm constraint."""
        lambda_max = self.mmep_cfg["lambda_max"]
        with torch.no_grad():
            for name, param in self.model.named_parameters():
                if param.dim() >= 2 and "weight" in name:
                    try:
                        s = torch.linalg.svdvals(param)
                        if s[0] > lambda_max:
                            param.data *= lambda_max / s[0]
                    except RuntimeError:
                        pass  # Skip non-matrix params

    def train_epoch(self, epoch: int) -> Dict[str, float]:
        """Iterate batches, run MMEP, update weights."""
        self.model.train()
        batch_size = self.train_cfg["batch_size"]
        hidden_dim = self.model_cfg["hidden_dim"]
        seq_len = min(self.data_cfg["max_seq_len"], 32)  # Cap for memory

        epoch_loss = 0.0
        epoch_entropy = 0.0
        num_batches = self.data_cfg["train_queries"] // batch_size

        for batch_idx in range(num_batches):
            # Synthetic arithmetic data
            x = torch.randn(batch_size, seq_len, hidden_dim, device=self.device)
            target = torch.randn(batch_size, seq_len, hidden_dim, device=self.device)

            self.optimizer.zero_grad()

            if self.use_mmep:
                # MMEP training: free phase then nudge phase
                h_free = self.mmep_free_phase(x)
                h_nudged = self.mmep_nudge_phase(x, target)

                # Local update signal
                update_signal = self.compute_local_update(h_free, h_nudged)

                # Use MSE of free-phase output vs target as differentiable loss
                out_free, aux_free = self.model(x)
                loss = F.mse_loss(out_free, target)

                # Add load balancing loss
                loss = loss + 0.01 * aux_free["lb_loss"]
            else:
                # Standard backpropagation
                output, aux = self.model(x)
                loss = F.mse_loss(output, target)
                loss = loss + 0.01 * aux["lb_loss"]
                aux_free = aux

            loss.backward()

            # Gradient clipping
            torch.nn.utils.clip_grad_norm_(
                self.model.parameters(),
                self.train_cfg["gradient_clip"]
            )

            self.optimizer.step()

            # Constraint projection (spectral norm)
            if self.use_constraint:
                self._spectral_norm_project()

            # Zero out memory parameters if ablation disables memory
            if not self.use_memory:
                with torch.no_grad():
                    for name, param in self.model.named_parameters():
                        if "memory" in name:
                            param.zero_()

            epoch_loss += loss.item()
            epoch_entropy += aux_free["entropy"].mean().item()

        avg_loss = epoch_loss / num_batches
        avg_entropy = epoch_entropy / num_batches

        return {
            "loss": avg_loss,
            "entropy": avg_entropy,
        }

    def evaluate(self, num_batches: int = 10) -> Dict[str, float]:
        """Run evaluation on synthetic data."""
        self.model.eval()
        batch_size = self.train_cfg["batch_size"]
        hidden_dim = self.model_cfg["hidden_dim"]
        seq_len = min(self.data_cfg["max_seq_len"], 32)

        total_loss = 0.0
        total_entropy = 0.0
        max_spectral = 0.0

        with torch.no_grad():
            for _ in range(num_batches):
                x = torch.randn(batch_size, seq_len, hidden_dim, device=self.device)
                target = torch.randn(batch_size, seq_len, hidden_dim, device=self.device)

                output, aux = self.model(x)
                loss = F.mse_loss(output, target)

                total_loss += loss.item()
                total_entropy += aux["entropy"].mean().item()

                # Track spectral norm
                for name, param in self.model.named_parameters():
                    if param.dim() >= 2 and "weight" in name:
                        try:
                            s = torch.linalg.svdvals(param)
                            max_spectral = max(max_spectral, s[0].item())
                        except RuntimeError:
                            pass

        return {
            "val_loss": total_loss / num_batches,
            "val_entropy": total_entropy / num_batches,
            "max_spectral_norm": max_spectral,
        }

    def run_ablation(self) -> Dict[str, List[float]]:
        """Run a specific ablation variant to completion."""
        print(f"\n{'='*60}")
        print(f"  Ablation: {self.ablation_name}")
        print(f"  MMEP: {self.use_mmep} | Memory: {self.use_memory} | "
              f"Constraint: {self.use_constraint} | MoE: {self.use_moe}")
        print(f"{'='*60}")

        num_epochs = self.train_cfg["num_epochs"]
        eval_every = self.train_cfg["eval_every"]
        checkpoint_every = self.train_cfg["checkpoint_every"]

        for epoch in range(1, num_epochs + 1):
            start_time = time.time()
            train_metrics = self.train_epoch(epoch)
            elapsed = time.time() - start_time

            metrics = {"epoch": epoch, "time_s": elapsed, **train_metrics}

            # Evaluation
            if epoch % eval_every == 0:
                val_metrics = self.evaluate()
                metrics.update(val_metrics)
                print(f"  Epoch {epoch:3d}/{num_epochs} | "
                      f"loss={train_metrics['loss']:.4f} | "
                      f"val_loss={val_metrics['val_loss']:.4f} | "
                      f"entropy={train_metrics['entropy']:.4f} | "
                      f"spectral={val_metrics['max_spectral_norm']:.4f} | "
                      f"time={elapsed:.1f}s")
            else:
                print(f"  Epoch {epoch:3d}/{num_epochs} | "
                      f"loss={train_metrics['loss']:.4f} | "
                      f"entropy={train_metrics['entropy']:.4f} | "
                      f"time={elapsed:.1f}s")

            self.metrics_history.append(metrics)

            # Save checkpoint
            if epoch % checkpoint_every == 0:
                self._save_checkpoint(epoch)

        # Save final metrics
        self._save_metrics()
        return self.metrics_history

    def run_all(self) -> Dict[str, List[Dict[str, float]]]:
        """Run all ablation variants defined in config."""
        all_results = {}
        ablation_names = [a["name"] for a in self.config.get("ablations", [])]

        for abl_name in ablation_names:
            # Create a fresh trainer for each ablation
            trainer = AblationTrainer(
                self.config,
                ablation_name=abl_name,
                device=str(self.device),
                output_dir=str(self.output_dir.parent),
            )
            results = trainer.run_ablation()
            all_results[abl_name] = results

        # Save comparative summary
        summary_path = self.output_dir.parent / "ablation_summary.json"
        summary = {}
        for name, history in all_results.items():
            if history:
                final = history[-1]
                summary[name] = {
                    "final_loss": final.get("loss", float("nan")),
                    "final_entropy": final.get("entropy", float("nan")),
                    "total_time": sum(h.get("time_s", 0) for h in history),
                }
        with open(summary_path, "w") as f:
            json.dump(summary, f, indent=2)
        print(f"\nAblation summary saved to: {summary_path}")

        return all_results

    def _save_checkpoint(self, epoch: int) -> None:
        """Save model checkpoint."""
        ckpt_path = self.output_dir / f"checkpoint_epoch_{epoch}.pt"
        torch.save({
            "epoch": epoch,
            "model_state_dict": self.model.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "config": self.config,
            "ablation": self.ablation_name,
        }, ckpt_path)

    def _save_metrics(self) -> None:
        """Save metrics history to JSON."""
        metrics_path = self.output_dir / "metrics.json"
        with open(metrics_path, "w") as f:
            json.dump(self.metrics_history, f, indent=2)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main():
    parser = argparse.ArgumentParser(
        description="BURT-IMMA Ablation Training Script"
    )
    parser.add_argument(
        "--config", type=str, default="config/ablation_arithmetic.yaml",
        help="Path to config YAML file"
    )
    parser.add_argument(
        "--ablation", type=str, default="full_mmep",
        help="Ablation variant name, or 'all' to run all variants"
    )
    parser.add_argument(
        "--device", type=str, default=None,
        help="Device (cpu, cuda:0, etc). Overrides config."
    )
    parser.add_argument(
        "--output-dir", type=str, default="results",
        help="Directory for checkpoints and metrics"
    )
    args = parser.parse_args()

    # Load config
    config_path = Path(args.config)
    if not config_path.exists():
        print(f"ERROR: Config file not found: {config_path}")
        raise SystemExit(1)

    with open(config_path) as f:
        config = yaml.safe_load(f)

    # Override device if specified
    device = args.device or config["experiment"].get("device", "cpu")
    if device.startswith("cuda") and not torch.cuda.is_available():
        print("WARNING: CUDA not available, falling back to CPU")
        device = "cpu"

    print(f"BURT-IMMA Training")
    print(f"  Config: {args.config}")
    print(f"  Device: {device}")
    print(f"  Output: {args.output_dir}")

    if args.ablation == "all":
        trainer = AblationTrainer(config, "full_mmep", device, args.output_dir)
        trainer.run_all()
    else:
        trainer = AblationTrainer(config, args.ablation, device, args.output_dir)
        trainer.run_ablation()

    print("\nTraining complete.")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""
	BURT-IMMA Full Training Script with Ablation Support
	License: BSL-1.1
	Contact: jessica@collectivekitty.com

	Loads config from config/ablation_arithmetic.yaml and runs MMEP training
	with ablation variants.

	Usage:
	python train_ablation.py --config config/ablation_arithmetic.yaml
	python train_ablation.py --ablation no_memory --device cuda:0
	python train_ablation.py --ablation all --output-dir results/
	"""

	import argparse
	import json
	import os
	import time
	from pathlib import Path
	from typing import Dict, List, Optional, Tuple

	import numpy as np
	import yaml

	try:
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	HAS_TORCH = True
	except ImportError:
	HAS_TORCH = False
	raise ImportError("PyTorch is required for training. Install with: pip install torch")


	# ---------------------------------------------------------------------------
	# Model Components
	# ---------------------------------------------------------------------------


	class SmoothLeakyActivation(nn.Module):
	"""alpha * x + (1-alpha) * x * sigmoid(x)"""

	def __init__(self, alpha: float = 0.01):
	super().__init__()
	self.alpha = alpha

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	return self.alpha * x + (1.0 - self.alpha) * x * torch.sigmoid(x)


	class GatesNormalization(nn.Module):
	"""Constrained softmax with entropy <= max_entropy."""

	def __init__(self, d_input: int, num_gates: int, max_entropy: float = 0.20):
	super().__init__()
	self.gate_proj = nn.Linear(d_input, num_gates)
	self.max_entropy = max_entropy

	def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
	logits = self.gate_proj(x)
	probs = F.softmax(logits, dim=-1)
	entropy = -(probs * (probs + 1e-10).log()).sum(dim=-1)

	# Sharpen via temperature reduction if entropy exceeds bound
	violations = entropy > self.max_entropy
	if violations.any():
	tau_lo = torch.full_like(entropy, 0.01)
	tau_hi = torch.ones_like(entropy)
	for _ in range(20):
	tau_mid = (tau_lo + tau_hi) / 2.0
	p = F.softmax(logits / tau_mid.unsqueeze(-1), dim=-1)
	h = -(p * (p + 1e-10).log()).sum(dim=-1)
	too_high = h > self.max_entropy
	tau_hi = torch.where(too_high, tau_mid, tau_hi)
	tau_lo = torch.where(too_high, tau_lo, tau_mid)
	tau_final = (tau_lo + tau_hi) / 2.0
	probs_sharp = F.softmax(logits / tau_final.unsqueeze(-1), dim=-1)
	probs = torch.where(violations.unsqueeze(-1), probs_sharp, probs)
	entropy = -(probs * (probs + 1e-10).log()).sum(dim=-1)

	return probs, entropy


	class CIFGMatrixMemory(nn.Module):
	"""Coupled Input-Forget Gate Matrix Memory."""

	def __init__(self, d_mem: int):
	super().__init__()
	self.d_mem = d_mem
	self.W_f = nn.Linear(d_mem, 1)
	self.register_buffer("C_global", torch.zeros(d_mem, d_mem))

	def forward(self, x: torch.Tensor, C_prev: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
	f = torch.sigmoid(self.W_f(x)) # (batch, 1)
	candidate = x.unsqueeze(-1) @ x.unsqueeze(-2) # (batch, d, d)
	candidate = candidate / (candidate.norm(dim=(-2, -1), keepdim=True) + 1e-8)
	C_new = f.unsqueeze(-1) * C_prev + (1.0 - f).unsqueeze(-1) * candidate
	return C_new, f

	def read(self, query: torch.Tensor, C: torch.Tensor) -> torch.Tensor:
	return torch.bmm(C, query.unsqueeze(-1)).squeeze(-1)


	class GatesRouter(nn.Module):
	"""Mixture-of-Experts router with top-k gating."""

	def __init__(self, d_input: int, num_experts: int, top_k: int = 1):
	super().__init__()
	self.num_experts = num_experts
	self.top_k = top_k
	self.gate = nn.Linear(d_input, num_experts)

	def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
	"""Returns (expert_indices, expert_weights, load_balance_loss)."""
	scores = self.gate(x)
	probs = F.softmax(scores, dim=-1)

	top_k_vals, top_k_idx = torch.topk(probs, self.top_k, dim=-1)
	top_k_vals = top_k_vals / (top_k_vals.sum(dim=-1, keepdim=True) + 1e-8)

	# Load balancing auxiliary loss
	routing_probs = probs.mean(dim=0)
	uniform = torch.ones(self.num_experts, device=x.device) / self.num_experts
	lb_loss = ((routing_probs - uniform) ** 2).sum()

	return top_k_idx, top_k_vals, lb_loss


	class SuperpositionedInductionHeads(nn.Module):
	"""Multi-head attention with induction bias."""

	def __init__(self, d_model: int, n_heads: int = 4):
	super().__init__()
	self.d_model = d_model
	self.n_heads = n_heads
	self.d_head = d_model // n_heads
	self.W_Q = nn.Linear(d_model, d_model)
	self.W_K = nn.Linear(d_model, d_model)
	self.W_V = nn.Linear(d_model, d_model)
	self.W_O = nn.Linear(d_model, d_model)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	batch, seq_len, d = x.shape
	Q = self.W_Q(x).view(batch, seq_len, self.n_heads, self.d_head).transpose(1, 2)
	K = self.W_K(x).view(batch, seq_len, self.n_heads, self.d_head).transpose(1, 2)
	V = self.W_V(x).view(batch, seq_len, self.n_heads, self.d_head).transpose(1, 2)

	scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_head ** 0.5)

	# Causal mask
	mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device) * float("-inf"), diagonal=1)
	scores = scores + mask.unsqueeze(0).unsqueeze(0)

	attn = F.softmax(scores, dim=-1)
	out = torch.matmul(attn, V)
	out = out.transpose(1, 2).contiguous().view(batch, seq_len, d)
	return self.W_O(out)


	class ExpertFFN(nn.Module):
	"""Single expert feed-forward network."""

	def __init__(self, d_model: int, d_ff: Optional[int] = None):
	super().__init__()
	d_ff = d_ff or d_model * 4
	self.W1 = nn.Linear(d_model, d_ff)
	self.W2 = nn.Linear(d_ff, d_model)
	self.activation = SmoothLeakyActivation()

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	return self.W2(self.activation(self.W1(x)))


	class BURTIMMALayer(nn.Module):
	"""Single BURT-IMMA layer."""

	def __init__(self, d_model: int, num_experts: int, top_k: int, d_mem: int):
	super().__init__()
	self.attention = SuperpositionedInductionHeads(d_model)
	self.norm1 = nn.LayerNorm(d_model)
	self.norm2 = nn.LayerNorm(d_model)
	self.router = GatesRouter(d_model, num_experts, top_k)
	self.experts = nn.ModuleList([ExpertFFN(d_model) for _ in range(num_experts)])
	self.memory = CIFGMatrixMemory(d_mem)

	def forward(self, x: torch.Tensor, C_prev: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
	# Attention with residual
	h = self.norm1(x + self.attention(x))

	# MoE routing
	batch, seq_len, d = h.shape
	h_flat = h.view(batch * seq_len, d)
	indices, weights, lb_loss = self.router(h_flat)

	# Dispatch to experts
	expert_out = torch.zeros_like(h_flat)
	for k in range(indices.shape[1]):
	for e_id in range(len(self.experts)):
	mask = (indices[:, k] == e_id)
	if mask.any():
	out_e = self.experts[e_id](h_flat[mask])
	expert_out[mask] += weights[mask, k:k+1] * out_e

	expert_out = expert_out.view(batch, seq_len, d)
	h = self.norm2(h + expert_out)

	# Memory update (use mean across sequence for memory key)
	h_mean = h.mean(dim=1) # (batch, d)
	C_new, f_gate = self.memory(h_mean, C_prev)

	return h, C_new, lb_loss


	class BURTIMMA(nn.Module):
	"""Full BURT-IMMA model."""

	def __init__(self, hidden_dim: int, num_layers: int, num_experts: int,
	top_k: int, d_mem: int):
	super().__init__()
	self.hidden_dim = hidden_dim
	self.num_layers = num_layers
	self.d_mem = d_mem

	self.input_proj = nn.Linear(hidden_dim, hidden_dim)
	self.layers = nn.ModuleList([
	BURTIMMALayer(hidden_dim, num_experts, top_k, d_mem)
	for _ in range(num_layers)
	])
	self.gates_norm = GatesNormalization(hidden_dim, num_experts)
	self.output_proj = nn.Linear(hidden_dim, hidden_dim)

	def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
	batch = x.shape[0]
	h = self.input_proj(x)

	C = torch.zeros(batch, self.d_mem, self.d_mem, device=x.device)
	total_lb_loss = torch.tensor(0.0, device=x.device)

	for layer in self.layers:
	h, C, lb_loss = layer(h, C)
	total_lb_loss = total_lb_loss + lb_loss

	# Gates normalization on final hidden
	h_mean = h.mean(dim=1)
	gate_probs, entropy = self.gates_norm(h_mean)

	out = self.output_proj(h)

	aux = {
	"entropy": entropy,
	"gate_probs": gate_probs,
	"lb_loss": total_lb_loss,
	"memory_trace": C.diagonal(dim1=-2, dim2=-1).sum(dim=-1),
	}
	return out, aux


	# ---------------------------------------------------------------------------
	# Ablation Trainer
	# ---------------------------------------------------------------------------


	class AblationTrainer:
	"""Full ablation training with MMEP support."""

	def __init__(self, config: dict, ablation_name: str = "full_mmep",
	device: str = "cpu", output_dir: str = "results"):
	self.config = config
	self.ablation_name = ablation_name
	self.device = torch.device(device)
	self.output_dir = Path(output_dir) / ablation_name
	self.output_dir.mkdir(parents=True, exist_ok=True)

	# Extract config sections
	self.model_cfg = config["model"]
	self.mmep_cfg = config["mmep"]
	self.train_cfg = config["training"]
	self.data_cfg = config["data"]

	# Apply ablation overrides
	self.use_mmep = True
	self.use_memory = True
	self.use_constraint = True
	self.use_moe = True
	self._apply_ablation(ablation_name)

	# Build model and optimizer
	self.model = self.build_model()
	self.model.to(self.device)
	self.optimizer = torch.optim.Adam(
	self.model.parameters(),
	lr=self.mmep_cfg["lr_W"]
	)

	# Metrics log
	self.metrics_history: List[Dict[str, float]] = []

	def _apply_ablation(self, name: str) -> None:
	"""Apply ablation variant configuration."""
	for ablation in self.config.get("ablations", []):
	if ablation["name"] == name:
	disable = ablation.get("disable", [])
	if "mmep" in disable:
	self.use_mmep = False
	if "C_global" in disable or "C_expert" in disable:
	self.use_memory = False
	if "projection" in disable:
	self.use_constraint = False
	if ablation.get("use_backprop", False):
	self.use_mmep = False
	if "override" in ablation:
	for k, v in ablation["override"].items():
	self.model_cfg[k] = v
	break

	def build_model(self) -> BURTIMMA:
	"""Create BURT-IMMA model from config."""
	num_experts = self.model_cfg["num_experts"] if self.use_moe else 1
	return BURTIMMA(
	hidden_dim=self.model_cfg["hidden_dim"],
	num_layers=self.model_cfg["num_layers"],
	num_experts=num_experts,
	top_k=self.model_cfg["top_k"],
	d_mem=self.model_cfg["d_mem"],
	)

	def mmep_free_phase(self, x: torch.Tensor) -> torch.Tensor:
	"""Run T_free relaxation steps (free phase of MMEP)."""
	T_free = self.mmep_cfg["T_free"]
	alpha = self.mmep_cfg["alpha_relax"]

	h = x.clone()
	for t in range(T_free):
	out, aux = self.model(h)
	# Relaxation: blend output back toward input
	h = (1.0 - alpha) * h + alpha * out

	return h

	def mmep_nudge_phase(self, x: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
	"""Run T_nudge nudged steps (nudged phase of MMEP)."""
	T_nudge = self.mmep_cfg["T_nudge"]
	beta = self.mmep_cfg["beta"]
	alpha = self.mmep_cfg["alpha_relax"]

	h = x.clone()
	for t in range(T_nudge):
	out, aux = self.model(h)
	h = (1.0 - alpha) * h + alpha * out
	# Nudge toward target
	h = h + beta * (target - h)

	return h

	def compute_local_update(self, h_free: torch.Tensor,
	h_nudged: torch.Tensor) -> torch.Tensor:
	"""Compute local parameter update: (nudged - free) / beta."""
	beta = self.mmep_cfg["beta"]
	return (h_nudged - h_free) / beta

	def _spectral_norm_project(self) -> None:
	"""Project weight matrices to satisfy spectral norm constraint."""
	lambda_max = self.mmep_cfg["lambda_max"]
	with torch.no_grad():
	for name, param in self.model.named_parameters():
	if param.dim() >= 2 and "weight" in name:
	try:
	s = torch.linalg.svdvals(param)
	if s[0] > lambda_max:
	param.data *= lambda_max / s[0]
	except RuntimeError:
	pass # Skip non-matrix params

	def train_epoch(self, epoch: int) -> Dict[str, float]:
	"""Iterate batches, run MMEP, update weights."""
	self.model.train()
	batch_size = self.train_cfg["batch_size"]
	hidden_dim = self.model_cfg["hidden_dim"]
	seq_len = min(self.data_cfg["max_seq_len"], 32) # Cap for memory

	epoch_loss = 0.0
	epoch_entropy = 0.0
	num_batches = self.data_cfg["train_queries"] // batch_size

	for batch_idx in range(num_batches):
	# Synthetic arithmetic data
	x = torch.randn(batch_size, seq_len, hidden_dim, device=self.device)
	target = torch.randn(batch_size, seq_len, hidden_dim, device=self.device)

	self.optimizer.zero_grad()

	if self.use_mmep:
	# MMEP training: free phase then nudge phase
	h_free = self.mmep_free_phase(x)
	h_nudged = self.mmep_nudge_phase(x, target)

	# Local update signal
	update_signal = self.compute_local_update(h_free, h_nudged)

	# Use MSE of free-phase output vs target as differentiable loss
	out_free, aux_free = self.model(x)
	loss = F.mse_loss(out_free, target)

	# Add load balancing loss
	loss = loss + 0.01 * aux_free["lb_loss"]
	else:
	# Standard backpropagation
	output, aux = self.model(x)
	loss = F.mse_loss(output, target)
	loss = loss + 0.01 * aux["lb_loss"]
	aux_free = aux

	loss.backward()

	# Gradient clipping
	torch.nn.utils.clip_grad_norm_(
	self.model.parameters(),
	self.train_cfg["gradient_clip"]
	)

	self.optimizer.step()

	# Constraint projection (spectral norm)
	if self.use_constraint:
	self._spectral_norm_project()

	# Zero out memory parameters if ablation disables memory
	if not self.use_memory:
	with torch.no_grad():
	for name, param in self.model.named_parameters():
	if "memory" in name:
	param.zero_()

	epoch_loss += loss.item()
	epoch_entropy += aux_free["entropy"].mean().item()

	avg_loss = epoch_loss / num_batches
	avg_entropy = epoch_entropy / num_batches

	return {
	"loss": avg_loss,
	"entropy": avg_entropy,
	}

	def evaluate(self, num_batches: int = 10) -> Dict[str, float]:
	"""Run evaluation on synthetic data."""
	self.model.eval()
	batch_size = self.train_cfg["batch_size"]
	hidden_dim = self.model_cfg["hidden_dim"]
	seq_len = min(self.data_cfg["max_seq_len"], 32)

	total_loss = 0.0
	total_entropy = 0.0
	max_spectral = 0.0

	with torch.no_grad():
	for _ in range(num_batches):
	x = torch.randn(batch_size, seq_len, hidden_dim, device=self.device)
	target = torch.randn(batch_size, seq_len, hidden_dim, device=self.device)

	output, aux = self.model(x)
	loss = F.mse_loss(output, target)

	total_loss += loss.item()
	total_entropy += aux["entropy"].mean().item()

	# Track spectral norm
	for name, param in self.model.named_parameters():
	if param.dim() >= 2 and "weight" in name:
	try:
	s = torch.linalg.svdvals(param)
	max_spectral = max(max_spectral, s[0].item())
	except RuntimeError:
	pass

	return {
	"val_loss": total_loss / num_batches,
	"val_entropy": total_entropy / num_batches,
	"max_spectral_norm": max_spectral,
	}

	def run_ablation(self) -> Dict[str, List[float]]:
	"""Run a specific ablation variant to completion."""
	print(f"\n{'='*60}")
	print(f" Ablation: {self.ablation_name}")
	print(f" MMEP: {self.use_mmep} \| Memory: {self.use_memory} \| "
	f"Constraint: {self.use_constraint} \| MoE: {self.use_moe}")
	print(f"{'='*60}")

	num_epochs = self.train_cfg["num_epochs"]
	eval_every = self.train_cfg["eval_every"]
	checkpoint_every = self.train_cfg["checkpoint_every"]

	for epoch in range(1, num_epochs + 1):
	start_time = time.time()
	train_metrics = self.train_epoch(epoch)
	elapsed = time.time() - start_time

	metrics = {"epoch": epoch, "time_s": elapsed, **train_metrics}

	# Evaluation
	if epoch % eval_every == 0:
	val_metrics = self.evaluate()
	metrics.update(val_metrics)
	print(f" Epoch {epoch:3d}/{num_epochs} \| "
	f"loss={train_metrics['loss']:.4f} \| "
	f"val_loss={val_metrics['val_loss']:.4f} \| "
	f"entropy={train_metrics['entropy']:.4f} \| "
	f"spectral={val_metrics['max_spectral_norm']:.4f} \| "
	f"time={elapsed:.1f}s")
	else:
	print(f" Epoch {epoch:3d}/{num_epochs} \| "
	f"loss={train_metrics['loss']:.4f} \| "
	f"entropy={train_metrics['entropy']:.4f} \| "
	f"time={elapsed:.1f}s")

	self.metrics_history.append(metrics)

	# Save checkpoint
	if epoch % checkpoint_every == 0:
	self._save_checkpoint(epoch)

	# Save final metrics
	self._save_metrics()
	return self.metrics_history

	def run_all(self) -> Dict[str, List[Dict[str, float]]]:
	"""Run all ablation variants defined in config."""
	all_results = {}
	ablation_names = [a["name"] for a in self.config.get("ablations", [])]

	for abl_name in ablation_names:
	# Create a fresh trainer for each ablation
	trainer = AblationTrainer(
	self.config,
	ablation_name=abl_name,
	device=str(self.device),
	output_dir=str(self.output_dir.parent),
	)
	results = trainer.run_ablation()
	all_results[abl_name] = results

	# Save comparative summary
	summary_path = self.output_dir.parent / "ablation_summary.json"
	summary = {}
	for name, history in all_results.items():
	if history:
	final = history[-1]
	summary[name] = {
	"final_loss": final.get("loss", float("nan")),
	"final_entropy": final.get("entropy", float("nan")),
	"total_time": sum(h.get("time_s", 0) for h in history),
	}
	with open(summary_path, "w") as f:
	json.dump(summary, f, indent=2)
	print(f"\nAblation summary saved to: {summary_path}")

	return all_results

	def _save_checkpoint(self, epoch: int) -> None:
	"""Save model checkpoint."""
	ckpt_path = self.output_dir / f"checkpoint_epoch_{epoch}.pt"
	torch.save({
	"epoch": epoch,
	"model_state_dict": self.model.state_dict(),
	"optimizer_state_dict": self.optimizer.state_dict(),
	"config": self.config,
	"ablation": self.ablation_name,
	}, ckpt_path)

	def _save_metrics(self) -> None:
	"""Save metrics history to JSON."""
	metrics_path = self.output_dir / "metrics.json"
	with open(metrics_path, "w") as f:
	json.dump(self.metrics_history, f, indent=2)


	# ---------------------------------------------------------------------------
	# Main
	# ---------------------------------------------------------------------------


	def main():
	parser = argparse.ArgumentParser(
	description="BURT-IMMA Ablation Training Script"
	)
	parser.add_argument(
	"--config", type=str, default="config/ablation_arithmetic.yaml",
	help="Path to config YAML file"
	)
	parser.add_argument(
	"--ablation", type=str, default="full_mmep",
	help="Ablation variant name, or 'all' to run all variants"
	)
	parser.add_argument(
	"--device", type=str, default=None,
	help="Device (cpu, cuda:0, etc). Overrides config."
	)
	parser.add_argument(
	"--output-dir", type=str, default="results",
	help="Directory for checkpoints and metrics"
	)
	args = parser.parse_args()

	# Load config
	config_path = Path(args.config)
	if not config_path.exists():
	print(f"ERROR: Config file not found: {config_path}")
	raise SystemExit(1)

	with open(config_path) as f:
	config = yaml.safe_load(f)

	# Override device if specified
	device = args.device or config["experiment"].get("device", "cpu")
	if device.startswith("cuda") and not torch.cuda.is_available():
	print("WARNING: CUDA not available, falling back to CPU")
	device = "cpu"

	print(f"BURT-IMMA Training")
	print(f" Config: {args.config}")
	print(f" Device: {device}")
	print(f" Output: {args.output_dir}")

	if args.ablation == "all":
	trainer = AblationTrainer(config, "full_mmep", device, args.output_dir)
	trainer.run_all()
	else:
	trainer = AblationTrainer(config, args.ablation, device, args.output_dir)
	trainer.run_ablation()

	print("\nTraining complete.")


	if __name__ == "__main__":
	main()