Rename experiments/exp_007_aleph_routed_attention/aleph_routed_trigram_lm.py to experiments/exp_007_aleph_routed_attention/2_aleph_routed_trigram_lm.py

experiments/exp_007_aleph_routed_attention/2_aleph_routed_trigram_lm.py

@@ -0,0 +1,537 @@
+# aleph_trigram_lm.py
+"""
+Aleph-Routed Trigram LM — the basin-test vehicle
+=================================================
+
+Causal language model over the SAME substrate the geolip-svae aleph batteries
+were fed: WikiText-103 as a raw UTF-8 byte stream, trigram-packed (3 bytes per
+position, stride 3 — the sequence analogue of ByteTrigramDataset's 3-bytes-per-
+cell RGB encoding). One token position = one trigram, embedded byte-factored
+(sum of 3 byte embeddings + position) and predicted as 3 independent 256-way
+byte heads — exactly the substrate, no 256^3 softmax.
+
+Purpose (preregistered — CORRECTED per RESEARCH_HISTORY.md):
+  The program established (Phase 2, discoveries #13/#15) at least TWO stable
+  codebook statutes, SELECTED BY SUBSTRATE: uniform-class (|dev| < 0.05; the
+  noise solvers) and polytope-class (dev > +0.05, pair fraction >= 45%;
+  repulsive packing — the BYTE-TRIGRAM solvers, in-distribution dev +0.083).
+  dev < -0.05 is degenerate (clumping) — the failure statute. Statute is a
+  property of (model x calibration), not the model alone.
+
+  Therefore, on THIS substrate, the basin question is statute-resolved:
+    - random init  -> STATUTE-SELECTION test. Codebook settling polytope-class
+      (the substrate-matched statute) or uniform-class under pure attention
+      gradients = cross-objective attractor evidence. Degenerate = failure.
+    - fibonacci init -> DRIFT test. Starts in the uniform basin; migration
+      OUT toward polytope under the symbolic substrate = substrate-driven
+      statute selection in a new objective (the stronger result).
+  Run BOTH inits. Statute (deviation + pair fraction, computed per the
+  program's own Sec 3.11 definitions) is logged per snapshot inline below;
+  void-richness (beta_2/axis via ripser on projective angular distances,
+  the symbolic-substrate fingerprint, discovery #20) is the deeper follow-up
+  on the saved snapshots. Note the two non-interchangeable "margins"
+  (top-1 softmax probability vs projective |<m,a>|): the inline monitors
+  detect collapse only; structure evidence is deviation/statute/beta_2.
+
+Substrate fidelity (mirrors geolip_svae.dataset_presets.ByteTrigramDataset):
+  - corpus: 'wikitext-103-raw-v1' via the Salesforce/wikitext namespace
+    (the bare 'wikitext' hub name is deprecated), or any local .txt path
+  - raw bytes in a single uint8 numpy array (prototypes/CLAUDE.md trap #3:
+    Python lists of ints are 5-7x the memory; never materialize them)
+  - seeded window sampling over the stream
+Repo invariants honored: pure Adam (never AdamW), no BatchNorm/Dropout on the
+geometric path, grad clip = max(task_loss, 1.0).
+
+Usage (Colab / Blackwell):
+    from aleph_trigram_lm import TrigramLMConfig, train_trigram_lm
+    cfg = TrigramLMConfig(steps=10_000, device='cuda',
+                          attn_mode='hub', codebook_init='random')
+    result = train_trigram_lm(cfg)
+    # result['codebook_snapshots'] -> [(step, (K, D_addr) tensor), ...]
+    # also saved to cfg.snapshot_path for the extraction pass
+
+Baseline A/B: attn_mode='standard' runs the same LM with softmax attention.
+
+Author: AbstractPhil + Mirel
+Date: 2026-06-09
+License: MIT
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import time
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from aleph_routed_attention import (
+    AlephRoutedAttention, AlephAttentionConfig, StandardAttention,
+)
+
+
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+# Config
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+@dataclass
+class TrigramLMConfig:
+    """Everything for one basin run.
+
+    Substrate:
+        corpus_id:        HF dataset config name (default = the aleph batteries'
+                          corpus) OR a local .txt/.text path
+        max_corpus_bytes: cap on bytes loaded (None = whole corpus, ~520 MB for
+                          wikitext-103). 50–100 MB is plenty for these runs.
+        seq_len:          context length in TRIGRAMS (bytes seen = 3*seq_len)
+
+    Model:
+        dim/n_layers/n_heads: transformer shell
+        attn_mode:        'hub' | 'bucket' | 'standard'
+        K/D_addr/tau:     aleph routing knobs (ignored for 'standard')
+        codebook_init:    'random' for the basin test (MANDATORY there) |
+                          'fibonacci' | (K, D_addr) array transplant
+
+    Training:
+        pure Adam + cosine decay to 10%; loss reported in nats and bits/byte.
+    """
+    # substrate
+    corpus_id: str = "wikitext-103-raw-v1"
+    split: str = "train"
+    max_corpus_bytes: Optional[int] = 100_000_000
+    seq_len: int = 256                      # trigrams (= 768 bytes of context)
+    seed: int = 1234
+
+    # model
+    dim: int = 384
+    n_layers: int = 4
+    n_heads: int = 6
+    attn_mode: str = "hub"                  # 'hub' | 'bucket' | 'standard'
+    K: int = 64
+    D_addr: int = 4
+    tau: float = 0.1
+    codebook_init: object = "random"        # basin test requires 'random'
+    div_weight: float = 0.0                 # anti-collapse; run 0 first, observe
+
+    # paradigm + scale
+    shared_codebook: bool = True            # ONE vocabulary, many speakers: all
+                                            # layers address the same (K,D) param,
+                                            # concentrating address pressure n_layers-x
+    accum_steps: int = 1                    # gradient accumulation (effective batch
+                                            # = batch_size * accum_steps)
+    stream_segments: int = 1                # segments per sample, each seq_len long;
+                                            # codebook-memory state carried across
+                                            # (TBPTT, detached between segments).
+                                            # context = seq_len * stream_segments
+                                            # at CONSTANT attention memory. hub-only.
+    probe_rows: int = 200_000               # address-stats sample size (ppl estimator)
+
+    # training
+    steps: int = 10_000                     # optimizer steps (micro-batches =
+                                            # steps * accum_steps)
+    batch_size: int = 32
+    lr: float = 3e-4
+    lr_decay: bool = True
+    log_every: int = 250
+    eval_batches: int = 8
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    amp: bool = False                       # bf16 autocast on the shell (the
+                                            # address stays fp32 inside)
+
+    # outputs
+    snapshot_codebook: bool = True
+    snapshot_path: str = "aleph_lm_codebook_snapshots.pt"
+    checkpoint_path: Optional[str] = "aleph_trigram_lm.pt"
+
+    def __post_init__(self):
+        assert self.attn_mode in ("hub", "bucket", "standard")
+        assert self.dim % self.n_heads == 0
+        assert self.accum_steps >= 1 and self.stream_segments >= 1
+        if self.stream_segments > 1:
+            assert self.attn_mode == "hub", \
+                "streaming requires mode='hub' (bucket sorts globally)"
+
+
+
+
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+# Statute monitor — the program's own diagnostic geometry (Sec 3.11)
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+def _canon(x: Tensor) -> Tensor:
+    """Sign-canonicalize onto RP^(D-1): flip so the first nonzero coord is
+    positive (antipodes map to one representative)."""
+    x = F.normalize(x, dim=-1)
+    first_nz = x[torch.arange(len(x)), x.abs().argmax(dim=-1)]
+    return x * torch.sign(first_nz).unsqueeze(-1)
+
+
+def _mean_projective_angle(X: Tensor) -> float:
+    """Mean pairwise acos|cos| over distinct pairs (radians)."""
+    c = (X @ X.t()).clamp(-1.0, 1.0).abs()
+    iu = torch.triu_indices(len(X), len(X), offset=1)
+    return torch.acos(c[iu[0], iu[1]]).mean().item()
+
+
+_UNIFORM_BASELINE: Dict[int, float] = {}
+
+def projective_deviation(axes: Tensor, n_ref: int = 4096,
+                         seed: int = 0) -> float:
+    """Uniformity deviation per the program definition: mean pairwise
+    projective angle of the axes MINUS the same statistic for n_ref uniform
+    random projective points at the same D. Signed; sign matters."""
+    D = axes.shape[-1]
+    if D not in _UNIFORM_BASELINE:
+        g = torch.Generator().manual_seed(seed)
+        ref = F.normalize(torch.randn(n_ref, D, generator=g), dim=-1)
+        _UNIFORM_BASELINE[D] = _mean_projective_angle(ref)
+    return _mean_projective_angle(F.normalize(axes.float(), dim=-1)) \
+        - _UNIFORM_BASELINE[D]
+
+
+def antipodal_pair_fraction(axes: Tensor, thresh: float = -0.9) -> float:
+    """Fraction of rows in mutual most-negative pairs with cos < thresh
+    (the antipodal-collapse acceptance rule)."""
+    A = F.normalize(axes.float(), dim=-1)
+    c = A @ A.t()
+    c.fill_diagonal_(0.0)
+    partner = c.argmin(dim=-1)
+    val = c.gather(-1, partner.unsqueeze(-1)).squeeze(-1)
+    mutual = partner[partner] == torch.arange(len(A))
+    return ((val < thresh) & mutual).float().mean().item()
+
+
+def statute(axes: Tensor) -> Dict[str, object]:
+    """Classify per the program taxonomy: dev > +0.05 polytope-class
+    (repulsive packing); |dev| < 0.05 uniform-class; dev < -0.05 degenerate
+    (clumping, the failure statute)."""
+    dev = projective_deviation(axes)
+    pf = antipodal_pair_fraction(axes)
+    cls = ("polytope" if dev > 0.05 else
+           "degenerate" if dev < -0.05 else "uniform")
+    return {"deviation": dev, "pair_fraction": pf, "statute": cls}
+
+
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+# Substrate — trigram stream (mirrors ByteTrigramDataset's loading)
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+class TrigramStream:
+    """WikiText-103 (or local .txt) as a uint8 byte stream, sampled as
+    causal trigram sequences.
+
+    __call__(batch, seq_len) -> (ids, targets):
+        ids:     (B, S, 3) uint8->long — trigram t  = bytes[3t : 3t+3]
+        targets: (B, S, 3)             — trigram t+1 (next-trigram prediction)
+    Windows are sampled at byte offsets aligned to stride 3 so the trigram
+    framing matches the image packing (cell i = bytes[3i : 3i+3])."""
+
+    def __init__(self, corpus_id: str, split: str = "train",
+                 max_corpus_bytes: Optional[int] = None, seed: int = 1234):
+        if os.path.isfile(corpus_id) and corpus_id.endswith((".txt", ".text")):
+            print(f"[TrigramStream] loading local corpus {corpus_id} ...")
+            with open(corpus_id, "rb") as f:
+                raw = f.read(max_corpus_bytes) if max_corpus_bytes else f.read()
+            self.stream = np.frombuffer(raw, dtype=np.uint8).copy()
+        else:
+            print(f"[TrigramStream] loading HF corpus {corpus_id} ...")
+            from datasets import load_dataset
+            if corpus_id.startswith("wikitext"):
+                ds = load_dataset("Salesforce/wikitext", corpus_id, split=split)
+            else:
+                ds = load_dataset(corpus_id, split=split)
+            # accumulate utf-8 bytes directly into a byte buffer — never a
+            # Python list of ints (prototypes/CLAUDE.md memory trap #3)
+            buf = bytearray()
+            cap = max_corpus_bytes or float("inf")
+            for row in ds:
+                t = row.get("text", "")
+                if t:
+                    buf.extend(t.encode("utf-8", errors="ignore"))
+                if len(buf) >= cap:
+                    break
+            self.stream = np.frombuffer(
+                bytes(buf[: max_corpus_bytes] if max_corpus_bytes else buf),
+                dtype=np.uint8).copy()
+        n_tri = len(self.stream) // 3
+        print(f"[TrigramStream] {len(self.stream):,} bytes "
+              f"= {n_tri:,} trigrams")
+        assert n_tri > 0, "corpus too small"
+        self._rng = np.random.default_rng(seed)
+
+    def sample(self, batch: int, seq_len: int,
+               device) -> Tuple[Tensor, Tensor]:
+        need = 3 * (seq_len + 1)                       # +1 trigram for targets
+        hi = len(self.stream) - need
+        assert hi > 0, f"corpus shorter than one window ({need} bytes)"
+        starts = self._rng.integers(0, hi // 3, size=batch) * 3   # stride-3 aligned
+        idx = starts[:, None] + np.arange(need)[None, :]          # (B, need)
+        window = self.stream[idx]                                  # (B, need) uint8
+        tri = torch.from_numpy(window.astype(np.int64)) \
+                   .view(batch, seq_len + 1, 3)
+        ids, targets = tri[:, :-1], tri[:, 1:]
+        return ids.to(device), targets.to(device)
+
+
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+# Model — byte-factored trigram LM
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+class TrigramLM(nn.Module):
+    """Causal LM over trigram positions. Embedding = sum of three per-slot
+    byte embeddings (+ learned positions); head = three 256-way byte heads.
+    Geometric-path hygiene: no BatchNorm/Dropout, pure pre-LN residual shell."""
+
+    def __init__(self, cfg: TrigramLMConfig):
+        super().__init__()
+        self.cfg = cfg
+        d = cfg.dim
+        self.byte_emb = nn.ModuleList([nn.Embedding(256, d) for _ in range(3)])
+        self.pos = nn.Parameter(0.02 * torch.randn(1, cfg.seq_len, d))
+
+        def make_attn() -> nn.Module:
+            if cfg.attn_mode == "standard":
+                return StandardAttention(d, cfg.n_heads, causal=True)
+            return AlephRoutedAttention(AlephAttentionConfig(
+                dim=d, num_heads=cfg.n_heads, mode=cfg.attn_mode,
+                K=cfg.K, D_addr=cfg.D_addr, tau=cfg.tau, causal=True,
+                codebook_init=cfg.codebook_init))
+
+        self.layers = nn.ModuleList([
+            nn.ModuleDict({
+                "norm1": nn.LayerNorm(d), "attn": make_attn(),
+                "norm2": nn.LayerNorm(d),
+                "mlp": nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(),
+                                     nn.Linear(4 * d, d)),
+            }) for _ in range(cfg.n_layers)
+        ])
+        self.norm_f = nn.LayerNorm(d)
+        self.heads = nn.ModuleList([nn.Linear(d, 256) for _ in range(3)])
+
+        # one vocabulary, many speakers: tie every layer's codebook to layer 0's
+        if cfg.shared_codebook and cfg.attn_mode in ("hub", "bucket"):
+            shared = self.layers[0]["attn"].codebook
+            for L in self.layers[1:]:
+                L["attn"].codebook = shared
+
+    def aleph_layers(self) -> List[AlephRoutedAttention]:
+        return [m for m in self.modules() if isinstance(m, AlephRoutedAttention)]
+
+    def backbone(self, ids: Tensor) -> Tensor:
+        """ids: (B, S, 3) -> (B, S, dim)"""
+        x = sum(emb(ids[..., i]) for i, emb in enumerate(self.byte_emb))
+        x = x + self.pos[:, : ids.shape[1]]
+        for L in self.layers:
+            x = x + L["attn"](L["norm1"](x))
+            x = x + L["mlp"](L["norm2"](x))
+        return self.norm_f(x)
+
+    def forward(self, ids: Tensor) -> Tensor:
+        """(B, S, 3) -> byte logits (B, S, 3, 256)"""
+        h = self.backbone(ids)
+        return torch.stack([head(h) for head in self.heads], dim=2)
+
+    def loss(self, ids: Tensor, targets: Tensor) -> Tensor:
+        """Mean cross-entropy per byte (nats/byte). bpb = loss / ln 2."""
+        logits = self(ids)                              # (B,S,3,256)
+        return F.cross_entropy(logits.reshape(-1, 256), targets.reshape(-1))
+
+    # ── streaming: segment-recurrent backbone (codebook memory carried) ──
+    def stream_loss(self, ids: Tensor, targets: Tensor,
+                    states: Optional[List] = None
+                    ) -> Tuple[Tensor, List]:
+        """One SEGMENT with per-layer carried states. states[i] is layer i's
+        (Mp, Mm, zp, zm) or None. Returns (loss, new_states) — caller detaches
+        between backward passes (TBPTT-1: grads flow within segment; values
+        flow forever)."""
+        x = sum(emb(ids[..., i]) for i, emb in enumerate(self.byte_emb))
+        x = x + self.pos[:, : ids.shape[1]]
+        new_states: List = []
+        states = states or [None] * len(self.layers)
+        for L, st in zip(self.layers, states):
+            a, ns = L["attn"].forward_stream(L["norm1"](x), state=st)
+            x = x + a
+            x = x + L["mlp"](L["norm2"](x))
+            new_states.append(ns)
+        h = self.norm_f(x)
+        logits = torch.stack([head(h) for head in self.heads], dim=2)
+        loss = F.cross_entropy(logits.reshape(-1, 256), targets.reshape(-1))
+        return loss, new_states
+
+
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+# Training
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+def train_trigram_lm(cfg: TrigramLMConfig,
+                     stream: Optional[TrigramStream] = None) -> Dict:
+    torch.manual_seed(cfg.seed)
+    dev = torch.device(cfg.device)
+    stream = stream or TrigramStream(cfg.corpus_id, cfg.split,
+                                     cfg.max_corpus_bytes, cfg.seed)
+    model = TrigramLM(cfg).to(dev)
+    n_params = sum(p.numel() for p in model.parameters())
+    opt = torch.optim.Adam(model.parameters(), lr=cfg.lr)   # pure Adam, never AdamW
+    sched = (torch.optim.lr_scheduler.CosineAnnealingLR(
+        opt, T_max=cfg.steps, eta_min=cfg.lr * 0.1) if cfg.lr_decay else None)
+
+    alephs = model.aleph_layers()
+    for a in alephs:
+        a.emit_diversity = cfg.div_weight > 0
+
+    snapshots: List[Tuple[int, Tensor]] = []
+    if cfg.snapshot_codebook and alephs:
+        snapshots.append((0, alephs[0].export_codebook()))
+
+    print(f"\n=== trigram LM  mode={cfg.attn_mode}  params={n_params:,}  "
+          f"ctx={cfg.seq_len}x{cfg.stream_segments} trigrams "
+          f"({3*cfg.seq_len*cfg.stream_segments} bytes)  "
+          f"eff.batch={cfg.batch_size*cfg.accum_steps}  "
+          f"shared_cb={cfg.shared_codebook}  "
+          f"device={dev} ===")
+    autocast = (torch.autocast(device_type=dev.type, dtype=torch.bfloat16)
+                if cfg.amp and dev.type == "cuda" else None)
+    result: Dict = {"mode": cfg.attn_mode, "params": n_params}
+    t0 = time.time()
+
+    segs = cfg.stream_segments
+    micro_scale = 1.0 / (cfg.accum_steps * segs)
+    for step in range(1, cfg.steps + 1):
+        opt.zero_grad(set_to_none=True)
+        loss_sum, n_micro = 0.0, 0
+        for _ in range(cfg.accum_steps):
+            # one long sample, split into `segs` carried segments
+            ids, targets = stream.sample(cfg.batch_size, cfg.seq_len * segs, dev)
+            states = None
+            for s in range(segs):
+                sl = slice(s * cfg.seq_len, (s + 1) * cfg.seq_len)
+                seg_ids, seg_tgt = ids[:, sl], targets[:, sl]
+                if autocast:
+                    with autocast:
+                        if segs > 1:
+                            loss, states = model.stream_loss(seg_ids, seg_tgt, states)
+                        else:
+                            loss = model.loss(seg_ids, seg_tgt)
+                else:
+                    if segs > 1:
+                        loss, states = model.stream_loss(seg_ids, seg_tgt, states)
+                    else:
+                        loss = model.loss(seg_ids, seg_tgt)
+                total = loss
+                if cfg.div_weight > 0 and alephs:
+                    total = total + cfg.div_weight * sum(
+                        a.diversity_loss() for a in alephs)
+                (total * micro_scale).backward()
+                loss_sum += loss.item(); n_micro += 1
+                if states is not None:                       # TBPTT boundary
+                    states = [tuple(t.detach() for t in st) for st in states]
+        loss_avg = loss_sum / n_micro
+        gnorm = torch.nn.utils.clip_grad_norm_(
+            model.parameters(), max(loss_avg, 1.0))           # the clip rule
+        opt.step()
+        if sched is not None:
+            sched.step()
+
+        if step % cfg.log_every == 0 or step == cfg.steps:
+            bpb = loss_avg / math.log(2)
+            rate = (step * cfg.batch_size * cfg.seq_len * segs
+                    * cfg.accum_steps) / (time.time() - t0)
+            line = (f"  step {step:6d}  loss {loss_avg:.4f}  "
+                    f"bpb {bpb:.3f}  |g| {gnorm:.2f}  {rate/1e3:.1f}k tri/s")
+            if alephs:
+                model.eval()
+                with torch.no_grad():
+                    x_probe = model.backbone(
+                        ids[: min(8, cfg.batch_size), -cfg.seq_len:])
+                st = alephs[0].address_stats(x_probe, max_rows=cfg.probe_rows)
+                model.train()
+                line += (f"  ppl {st['perplexity']:.1f}/{st['max_perplexity']:.0f}"
+                         f"  margin {st['margin']:.4f}"
+                         f"  conf {st['confidence']:.3f}")
+                result.update(st)
+                if cfg.snapshot_codebook:
+                    snapshots.append((step, alephs[0].export_codebook()))
+            print(line)
+            result.update({"loss": loss_avg, "bpb": bpb, "step": step})
+
+    if snapshots:
+        result["codebook_snapshots"] = snapshots
+        drift = (snapshots[-1][1] - snapshots[0][1]).norm().item()
+        result["codebook_drift"] = drift
+        traj = [(s, statute(cb)) for s, cb in snapshots]
+        result["statute_trajectory"] = traj
+        torch.save({"snapshots": snapshots, "statute_trajectory": traj,
+                    "config": cfg.__dict__, "K": cfg.K, "D_addr": cfg.D_addr},
+                   cfg.snapshot_path)
+        print(f"\n[basin] {len(snapshots)} snapshots -> {cfg.snapshot_path}"
+              f"   drift |A_end - A_0| = {drift:.4f}")
+        print("[basin] statute trajectory (program taxonomy: polytope is the "
+              "substrate-matched\n        statute for byte-trigram; uniform is "
+              "the noise/OOD statute; degenerate = failure):")
+        for s, st in traj:
+            print(f"    step {s:6d}  dev {st['deviation']:+.4f}  "
+                  f"pairs {st['pair_fraction']:.0%}  -> {st['statute']}")
+        print("[basin] deeper follow-up on saved snapshots: beta_2/axis via "
+              "ripser on projective\n        angular distances (the "
+              "void/symbolic fingerprint, discovery #20).")
+    if cfg.checkpoint_path:
+        torch.save({"model_state_dict": model.state_dict(),
+                    "config": cfg.__dict__}, cfg.checkpoint_path)
+        print(f"[ckpt] saved -> {cfg.checkpoint_path}")
+    return result
+
+
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+# Activation
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+def _smoke(device: str = "cpu"):
+    """End-to-end on a synthetic local corpus — no downloads."""
+    print("=" * 70)
+    print("aleph_trigram_lm — smoke (synthetic corpus)")
+    print("=" * 70)
+    path = "/tmp/_smoke_corpus.txt"
+    rng = np.random.default_rng(0)
+    words = [b"the", b"aleph", b"address", b"routes", b"attention",
+             b"through", b"a", b"learned", b"projective", b"codebook"]
+    with open(path, "wb") as f:
+        f.write(b" ".join(words[i] for i in rng.integers(0, len(words), 60_000)))
+    cfg = TrigramLMConfig(corpus_id=path, steps=30, log_every=10,
+                          dim=96, n_layers=2, n_heads=4, K=16, seq_len=64,
+                          batch_size=8, device=device,
+                          snapshot_path="/tmp/_smoke_snaps.pt",
+                          checkpoint_path=None)
+    r = train_trigram_lm(cfg)
+    assert "codebook_snapshots" in r and len(r["codebook_snapshots"]) >= 2
+    assert math.isfinite(r["loss"]) and r["loss"] < math.log(256)
+    print(f"\nsmoke OK — loss {r['loss']:.3f} (< ln256={math.log(256):.3f} prior), "
+          f"drift {r['codebook_drift']:.4f}, "
+          f"{len(r['codebook_snapshots'])} snapshots saved")
+
+
+if __name__ == "__main__":
+    import argparse
+    ap = argparse.ArgumentParser(description="Aleph trigram LM — basin run")
+    ap.add_argument("--smoke-only", action="store_true")
+    ap.add_argument("--mode", default="hub", choices=["hub", "bucket", "standard"])
+    ap.add_argument("--steps", type=int, default=10_000)
+    ap.add_argument("--corpus-mb", type=int, default=100)
+    ap.add_argument("--codebook-init", default="random")
+    ap.add_argument("--device",
+                    default="cuda" if torch.cuda.is_available() else "cpu")
+    args, _unknown = ap.parse_known_args()      # notebook-safe (ignores -f kernel.json)
+
+    if args.smoke_only:
+        _smoke(device="cpu")
+    else:
+        cfg = TrigramLMConfig(attn_mode=args.mode, steps=args.steps,
+                              max_corpus_bytes=args.corpus_mb * 1_000_000,
+                              codebook_init=args.codebook_init,
+                              device=args.device)
+        train_trigram_lm(cfg)

experiments/exp_007_aleph_routed_attention/aleph_routed_trigram_lm.py

@@ -1,792 +0,0 @@
-# aleph_lm.py
-"""
-AlephLM — prediction through the codebook, with guarantees
-===========================================================
-
-The composite reduction (2026-06-09): a causal trigram LM in which the aleph
-signed-projective address is load-bearing at ALL THREE stations — input
-addressing, mixing, and prediction. The codebook receives gradient from
-routing, from the predicted next-address pi, and from every candidate address
-kappa. One geometry, closed loop, smooth everywhere (no argmax in the train
-path): differential trigram-to-trigram prediction.
-
-    CODEC      bytes -> trigrams g_t (stride 3)
-    EMBED      e_t = sum_c E_c[g_t[c]]                       (byte-factored)
-    MIX        AlephRoutedAttention hub layers, shared codebook, causal
-    PREDICT    pi = softmax([w; -w]),  w = W_pi h_t           (free antipodal-tied)
-    CANDIDATE  kappa(tau) = address(normalize(W_k sum_c E_c[tau[c]]))
-    SCORE      logit(tau) = alpha * log( pi+ . kappa+(tau) + pi- . kappa-(tau) )
-    OUTPUT     hybrid:  P(g) = g_in * P_bank(g | in) + (1-g_in) * P_byte(g)
-
-THE GUARANTEE LEDGER (all demonstrated numerically 2026-06-09; see session log):
-  T1/T2  pi parameterization: address-constrained pi is projectively UNIMODAL
-         (logits linear in x-hat) — a two-spike target is unreachable (best
-         joint mass 1.4% vs 50% needed). The free antipodal-tied simplex
-         represents any tied-logit distribution. DEFAULT: free tied simplex;
-         address-constrained is the unimodal ablation (pi_mode='address').
-  T3     Tied [w; -w] implies p+k * p-k is CONSTANT across k — every axis is
-         forced to an orientation stance. Feature-or-bug: empirical.
-  T4     The 3x256 byte-product head cannot express within-trigram byte
-         correlation (rank-1 tensor over 256^3); it is the guaranteed-floor
-         baseline (head='byte'), not the main head.
-  T5     The hybrid output is a PROPER full-support distribution and its CE
-         decomposes exactly: -log P(g) = -log gate_branch - log P_branch(g).
-         Implemented verbatim. "Run all three banks" = ablations inside one
-         provably-correct machine.
-  T6     Raw-score softmax over a bank has a sharpness ceiling (scores in
-         (0,1] => CE floor 7.32 nats at M=4096). Logits are LOG-kernel with a
-         learnable scale alpha. Non-negotiable.
-  T7     Output logit rank <= 2K (softmax bottleneck): K governs attention
-         rank, output rank, and mode capacity — one knob, three proven roles.
-  T8     The write-head target Delta-z = sum of future addresses is the
-         ORDER-MARGINALIZED multiset of the next W trigrams (permutation-
-         invariant by commutativity). It predicts WHAT comes, not the order.
-         Learnability rests on the empirical rank-10 occupancy result.
-  Lit.   Sampled softmax requires the log-Q correction; with a uniform
-         proposal the correction is constant and cancels in the softmax
-         (target always included). head='sampled' implements exactly this.
-
-THE BRANCHING GAUGE (the [TAU] kernel invariant, inverted): a single unit row
-has conf = ||(p+ - p-)A|| pinned at f(tau,K,D). A PREDICTED pi is not so
-bound — implied confidence below the invariant is the model declaring
-superposition. branching_frac is monitored from step zero.
-
-Banks: 'corpus' (top-M trigrams of the training stream), 'wordnet'
-(AbstractPhil/wordnet-lexical-topology char_eng_3gram, frequency-ranked,
-filtered to exact 3-byte UTF-8), or per-step 'sampled' negatives.
-
-Usage (Blackwell / A100):
-    from aleph_lm import AlephLMConfig, train_aleph_lm
-    r = train_aleph_lm(AlephLMConfig(steps=10_000, device='cuda',
-                                     head='hybrid', bank_source='wordnet'))
-
-Depends: aleph_routed_attention.py, aleph_trigram_lm.py in the same directory.
-Author: AbstractPhil + Mirel    Date: 2026-06-09    License: MIT
-"""
-
-from __future__ import annotations
-
-import math
-import os
-import time
-from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch import Tensor
-
-# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-# Config
-# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-
-@dataclass
-class AlephLMConfig:
-    # substrate
-    corpus_id: str = "wikitext-103-raw-v1"
-    split: str = "train"
-    max_corpus_bytes: Optional[int] = 100_000_000
-    seq_len: int = 256                       # trigrams (3*seq_len bytes)
-    seed: int = 1234
-
-    # tower
-    dim: int = 384
-    n_layers: int = 4
-    n_heads: int = 6
-    K: int = 64
-    D_addr: int = 4
-    tau: float = 0.1
-    codebook_init: object = "random"
-    shared_codebook: bool = True
-
-    # prediction head (the ledger's resolutions)
-    head: str = "hybrid"                     # 'hybrid' | 'byte' | 'bank' | 'sampled'
-    pi_mode: str = "free"                    # 'free' (T2 default) | 'address' (unimodal ablation)
-    bank_source: str = "corpus"              # 'corpus' | 'wordnet'
-    bank_size: int = 4096
-    n_negatives: int = 1024                  # head='sampled'
-    logit_scale_init: float = 1.0            # alpha on the log-kernel logits (T6)
-
-    # hybrid bank scorer: 'kernel' (log-kernel, T6) or 'pmix' — a mixture of
-    # J pointers on S^(d_point-1): logits(c) = logsumexp_j [log w_j + T yhat_j.c]
-    # = Mixture-of-Softmaxes in sphere coordinates. Theorem-backed twice over:
-    # raises output rank past the T7 bottleneck (MoS, Yang et al.), and gives
-    # the pointer J modes so the barycenter pathology (unimodal aim at a
-    # multimodal future) is structurally removed. Full-bank softmax retained:
-    # T5 propriety intact. PREREGISTERED statute prediction: pmix candidate
-    # coords bypass the codebook (W_cand48), removing prediction-side
-    # discrimination pressure -> expect dev near the zero group, vs kernel's
-    # +0.013. The dose-response gets a within-architecture test.
-    bank_scorer: str = "kernel"              # 'kernel' | 'pmix'
-    n_pointers: int = 4                      # J mixture components (pmix)
-
-    # pointer head (head='pointer'): NN-on-the-sphere decode
-    d_point: int = 48                        # pointer sphere dim (band-valid; the
-                                             # capacity table gives the decode
-                                             # budget theta_NN/2 at this D)
-    pointer_k: int = 32                      # hard negatives = target's k sphere-NN
-    pointer_cos_weight: float = 0.5          # aiming regularizer (contrastive CE
-                                             # is the main learner — lit. caveat)
-    pointer_refresh: int = 200               # steps between NN-table refreshes
-                                             # (candidate coords drift)
-
-    # write-head (T8, auxiliary multiset prediction)
-    write_weight: float = 0.1                # 0 disables
-    write_horizon: int = 8                   # W: the granularity dial
-
-    # training
-    steps: int = 10_000
-    batch_size: int = 32
-    accum_steps: int = 1
-    lr: float = 3e-4
-    lr_decay: bool = True
-    div_weight: float = 0.0
-    log_every: int = 250
-    device: str = "cuda" if torch.cuda.is_available() else "cpu"
-
-    # outputs
-    snapshot_codebook: bool = True
-    snapshot_path: str = "aleph_lm_snaps.pt"
-    checkpoint_path: Optional[str] = "aleph_lm.pt"
-
-    def __post_init__(self):
-        assert self.head in ("hybrid", "byte", "bank", "sampled", "pointer")
-        assert self.pi_mode in ("free", "address")
-        assert self.bank_scorer in ("kernel", "pmix")
-        assert self.bank_source in ("corpus", "wordnet") \
-            or os.path.isfile(str(self.bank_source)), \
-            f"bank_source must be 'corpus'|'wordnet'|path to bank .pt"
-        assert self.dim % self.n_heads == 0
-        assert self.write_horizon >= 1
-        tag = self.head
-        if self.head in ("hybrid", "bank"):
-            b = (os.path.splitext(os.path.basename(str(self.bank_source)))[0]
-                 if os.path.isfile(str(self.bank_source)) else self.bank_source)
-            tag += f"_{b}"
-        if self.pi_mode != "free":
-            tag += f"_{self.pi_mode}"
-        if self.head == "hybrid" and self.bank_scorer == "pmix":
-            tag += f"_pmix{self.n_pointers}"
-        if self.checkpoint_path == "aleph_lm.pt":
-            self.checkpoint_path = f"aleph_lm_{tag}.pt"
-        if self.snapshot_path == "aleph_lm_snaps.pt":
-            self.snapshot_path = f"aleph_lm_snaps_{tag}.pt"
-
-
-# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-# Candidate banks
-# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-
-def _tri_ids(tri: Tensor) -> Tensor:
-    """(..., 3) bytes -> scalar trigram id in [0, 256^3)."""
-    return tri[..., 0] * 65536 + tri[..., 1] * 256 + tri[..., 2]
-
-
-def build_corpus_bank(stream: TrigramStream, M: int,
-                      sample_bytes: int = 6_000_000) -> Tensor:
-    """Top-M most frequent trigrams of the training stream. (M, 3) long."""
-    n = min(sample_bytes, (len(stream.stream) // 3) * 3)
-    tri = stream.stream[:n].reshape(-1, 3)
-    ids = (tri[:, 0].astype(np.int64) * 65536 + tri[:, 1].astype(np.int64) * 256
-           + tri[:, 2].astype(np.int64))
-    uniq, counts = np.unique(ids, return_counts=True)
-    top = uniq[np.argsort(counts)[::-1][:M]]
-    out = np.stack([top // 65536, (top // 256) % 256, top % 256], axis=-1)
-    return torch.from_numpy(out.astype(np.int64))
-
-
-def build_wordnet_bank(M: int) -> Tensor:
-    """char_eng_3gram from AbstractPhil/wordnet-lexical-topology, frequency-
-    ranked, filtered to exact 3-byte UTF-8. (M', 3) long, M' <= M."""
-    from huggingface_hub import hf_hub_download
-    import pyarrow.parquet as pq
-    p = hf_hub_download("AbstractPhil/wordnet-lexical-topology",
-                        "data/char_eng_3gram-00000-of-00001.parquet",
-                        repo_type="dataset")
-    t = pq.read_table(p, columns=["ngram", "rank"]).to_pandas()
-    t = t.sort_values("rank")
-    rows = []
-    for s in t["ngram"]:
-        b = str(s).encode("utf-8", errors="ignore")
-        if len(b) == 3:
-            rows.append([b[0], b[1], b[2]])
-        if len(rows) >= M:
-            break
-    assert rows, "wordnet bank empty after 3-byte filter"
-    return torch.tensor(rows, dtype=torch.long)
-
-
-# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-# Model
-# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-
-class AlephLM(nn.Module):
-    """The composite reduction. forward_loss(ids, targets) -> (loss, logs)."""
-
-    def __init__(self, cfg: AlephLMConfig, bank: Optional[Tensor] = None):
-        super().__init__()
-        self.cfg = cfg
-        d = cfg.dim
-
-        # ── EMBED (byte-factored; shared with candidate composition) ──
-        self.byte_emb = nn.ModuleList([nn.Embedding(256, d) for _ in range(3)])
-        self.pos = nn.Parameter(0.02 * torch.randn(1, cfg.seq_len, d))
-
-        # ── MIX (hub tower, shared codebook) ──
-        def make_attn():
-            return AlephRoutedAttention(AlephAttentionConfig(
-                dim=d, num_heads=cfg.n_heads, mode="hub", K=cfg.K,
-                D_addr=cfg.D_addr, tau=cfg.tau, causal=True,
-                codebook_init=cfg.codebook_init))
-        self.layers = nn.ModuleList([
-            nn.ModuleDict({"norm1": nn.LayerNorm(d), "attn": make_attn(),
-                           "norm2": nn.LayerNorm(d),
-                           "mlp": nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(),
-                                                nn.Linear(4 * d, d))})
-            for _ in range(cfg.n_layers)])
-        if cfg.shared_codebook:
-            shared = self.layers[0]["attn"].codebook
-            for L in self.layers[1:]:
-                L["attn"].codebook = shared
-        self.norm_f = nn.LayerNorm(d)
-
-        # ── PREDICT: pi over 2K oriented axes ──
-        self.W_pi = nn.Linear(d, cfg.K, bias=True)          # tied logits [w; -w]
-        if cfg.pi_mode == "address":                        # unimodal ablation (T2)
-            self.W_pi_row = nn.Linear(d, cfg.D_addr, bias=False)
-            nn.init.orthogonal_(self.W_pi_row.weight)
-
-        # ── CANDIDATE: compositional addresses (banked heads) ──
-        self.W_kappa = nn.Linear(d, cfg.D_addr, bias=False)
-        nn.init.orthogonal_(self.W_kappa.weight)
-        self.logit_scale = nn.Parameter(
-            torch.tensor(float(cfg.logit_scale_init)))      # alpha (T6)
-
-        # ── byte-product head (T4 floor + hybrid tail) ──
-        self.byte_heads = nn.ModuleList([nn.Linear(d, 256) for _ in range(3)])
-
-        # ── hybrid gate (T5) ──
-        self.gate = nn.Linear(d, 1)
-
-        # ── pmix bank scorer: J-pointer mixture (MoS on the sphere) ──
-        if cfg.head == "hybrid" and cfg.bank_scorer == "pmix":
-            J = cfg.n_pointers
-            self.W_pmix = nn.Linear(d, J * cfg.d_point, bias=False)
-            nn.init.orthogonal_(self.W_pmix.weight)
-            self.W_mixgate = nn.Linear(d, J)
-            self.W_cand48 = nn.Linear(d, cfg.d_point, bias=False)
-            nn.init.orthogonal_(self.W_cand48.weight)
-            self.point_T = nn.Parameter(torch.tensor(10.0))
-
-        # ── pointer head: predict a point on S^(d_point-1), decode by NN ──
-        if cfg.head == "pointer":
-            self.W_point = nn.Linear(d, cfg.d_point, bias=False)
-            nn.init.orthogonal_(self.W_point.weight)
-            self.W_cand48 = nn.Linear(d, cfg.d_point, bias=False)
-            nn.init.orthogonal_(self.W_cand48.weight)
-            self.point_T = nn.Parameter(torch.tensor(10.0))  # contrastive inv-temp
-            self.register_buffer("_nn_table", torch.zeros(0, dtype=torch.long),
-                                 persistent=False)
-            self._nn_step = -1
-
-        # ── write-head (T8): predicted Delta-z over 2K ──
-        if cfg.write_weight > 0:
-            self.W_write = nn.Linear(d, 2 * cfg.K)
-
-        # ── bank registration ──
-        if bank is not None:
-            self.register_buffer("bank", bank)                      # (M, 3)
-            self.register_buffer("bank_ids_sorted",
-                                 _tri_ids(bank).sort().values)      # membership
-            self.register_buffer("bank_perm",
-                                 _tri_ids(bank).argsort())          # sorted->orig
-        else:
-            self.bank = None
-
-    # ---- shared codebook handle ----
-    @property
-    def codebook(self) -> Tensor:
-        return self.layers[0]["attn"].codebook
-
-    def aleph_layers(self) -> List[AlephRoutedAttention]:
-        return [m for m in self.modules() if isinstance(m, AlephRoutedAttention)]
-
-    # ---- address of arbitrary unit rows vs the SHARED codebook ----
-    def _address_rows(self, rows: Tensor) -> Tuple[Tensor, Tensor]:
-        A = F.normalize(self.codebook, dim=-1)
-        u = (rows @ A.t()) / self.cfg.tau
-        m = u.abs().amax(-1, keepdim=True)
-        ep, en = torch.exp(u - m), torch.exp(-u - m)
-        Z = (ep + en).sum(-1, keepdim=True)
-        return ep / Z, en / Z
-
-    # ---- PREDICT ----
-    def _pi(self, h: Tensor) -> Tuple[Tensor, Tensor]:
-        """pi over 2K oriented axes. 'free': tied simplex (T2 default).
-        'address': unimodal ablation."""
-        if self.cfg.pi_mode == "address":
-            row = F.normalize(self.W_pi_row(h), dim=-1)
-            return self._address_rows(row)
-        w = self.W_pi(h)                                    # (..., K)
-        m = w.abs().amax(-1, keepdim=True)
-        ep, en = torch.exp(w - m), torch.exp(-w - m)
-        Z = (ep + en).sum(-1, keepdim=True)
-        return ep / Z, en / Z
-
-    # ---- CANDIDATE addresses for a (M, 3) byte bank ----
-    def _kappa(self, bank: Tensor) -> Tuple[Tensor, Tensor]:
-        e = sum(emb(bank[:, i]) for i, emb in enumerate(self.byte_emb))
-        rows = F.normalize(self.W_kappa(e), dim=-1)         # (M, D_addr)
-        return self._address_rows(rows)
-
-    # ---- SCORE: log-kernel logits (T6) ----
-    def _bank_logits(self, pi_p: Tensor, pi_m: Tensor,
-                     k_p: Tensor, k_m: Tensor) -> Tensor:
-        s = pi_p @ k_p.t() + pi_m @ k_m.t()                 # (..., M), > 0
-        return self.logit_scale * torch.log(s.clamp_min(1e-9))
-
-    # ---- pmix bank logits: logsumexp over J pointers (MoS on the sphere) ----
-    def _pmix_logits(self, h: Tensor) -> Tuple[Tensor, Dict[str, float]]:
-        cfg = self.cfg
-        J = cfg.n_pointers
-        coords = self._cand_coords()                          # (M, d_point)
-        y = self.W_pmix(h).view(*h.shape[:-1], J, cfg.d_point)
-        y = F.normalize(y, dim=-1)                            # (B,S,J,dp)
-        mix = F.log_softmax(self.W_mixgate(h), dim=-1)        # (B,S,J)
-        sims = torch.einsum("bsjd,md->bsjm", y, coords) * self.point_T
-        logits = torch.logsumexp(mix.unsqueeze(-1) + sims, dim=2)   # (B,S,M)
-        with torch.no_grad():                                 # mode diagnostics
-            pw = torch.einsum("bsjd,bskd->bsjk", y, y)
-            off = pw.masked_select(~torch.eye(J, dtype=torch.bool,
-                                              device=h.device)
-                                   .expand_as(pw)).clamp(-1, 1)
-            spread = torch.acos(off).mean().item() * 180 / math.pi
-            usage = mix.exp().mean(dim=(0, 1))
-            ent = -(usage * usage.clamp_min(1e-9).log()).sum().item() / math.log(J)
-        return logits, {"mode_spread_deg": spread, "mix_entropy": ent}
-
-    # ---- tower ----
-    def backbone(self, ids: Tensor) -> Tensor:
-        x = sum(emb(ids[..., i]) for i, emb in enumerate(self.byte_emb))
-        x = x + self.pos[:, : ids.shape[1]]
-        for L in self.layers:
-            x = x + L["attn"](L["norm1"](x))
-            x = x + L["mlp"](L["norm2"](x))
-        return self.norm_f(x)
-
-    # ---- byte-product log-probs of given targets (T4 tail) ----
-    def _byte_logprob(self, h: Tensor, targets: Tensor) -> Tensor:
-        lp = 0.0
-        for c, head in enumerate(self.byte_heads):
-            lp = lp + F.log_softmax(head(h), dim=-1).gather(
-                -1, targets[..., c:c + 1]).squeeze(-1)
-        return lp                                            # (B, S)
-
-    # ---- bank membership: target -> bank index or -1 ----
-    def _bank_index(self, targets: Tensor) -> Tensor:
-        tid = _tri_ids(targets)
-        pos = torch.searchsorted(self.bank_ids_sorted, tid)
-        pos = pos.clamp_max(len(self.bank_ids_sorted) - 1)
-        hit = self.bank_ids_sorted[pos] == tid
-        idx = self.bank_perm[pos]
-        return torch.where(hit, idx, torch.full_like(idx, -1))
-
-    # ---- write-head target (T8): order-marginalized future address mass ----
-    @torch.no_grad()
-    def _write_target(self, ids: Tensor) -> Tensor:
-        """Delta-z over 2K for horizon W at each position (normalized)."""
-        cfg = self.cfg
-        a0 = self.layers[0]["attn"]
-        x = sum(emb(ids[..., i]) for i, emb in enumerate(self.byte_emb))
-        x = x + self.pos[:, : ids.shape[1]]
-        kh = a0._split_addr(a0.k_addr(self.layers[0]["norm1"](x)),
-                            ids.shape[0], ids.shape[1])
-        pk_p, pk_m = a0._address(kh)                         # (B,H,S,K)
-        p = torch.cat([pk_p, pk_m], dim=-1).mean(dim=1)      # (B,S,2K)
-        cs = torch.cat([torch.zeros_like(p[:, :1]), p.cumsum(dim=1)], dim=1)
-        W = cfg.write_horizon
-        B, S, _ = p.shape
-        end = torch.arange(S, device=p.device).clamp_max(S - 1)
-        lo = cs[:, 1:]                                       # prefix up to t (incl)
-        hi = cs[:, torch.clamp(torch.arange(S, device=p.device) + W, max=S)]
-        dz = (hi - lo).clamp_min(0)
-        valid = (torch.arange(S, device=p.device) + 1 < S)   # at least 1 future tok
-        dz = dz / dz.sum(-1, keepdim=True).clamp_min(1e-9)
-        return dz, valid
-
-    # ---- pointer head: compositional D=48 candidate coordinates ----
-    def _cand_coords(self) -> Tensor:
-        e = sum(emb(self.bank[:, i]) for i, emb in enumerate(self.byte_emb))
-        return F.normalize(self.W_cand48(e), dim=-1)         # (M, d_point)
-
-    @torch.no_grad()
-    def _refresh_nn(self, coords: Tensor, step: int) -> None:
-        """Hard-negative table: each candidate's k nearest sphere neighbors
-        (excluding self). Refreshed periodically — coordinates drift."""
-        cos = coords @ coords.t()
-        cos.fill_diagonal_(-2.0)
-        self._nn_table = cos.topk(self.cfg.pointer_k, dim=-1).indices  # (M, k)
-        self._nn_step = step
-        # decode budget: theta_NN/2 of the CURRENT candidate constellation
-        nn_deg = torch.acos(cos.max(dim=-1).values.clamp(-1, 1)) * 180 / math.pi
-        self._decode_budget_deg = (nn_deg.median() / 2).item()
-
-    def _pointer_loss(self, h: Tensor, targets: Tensor,
-                      step: int) -> Tuple[Tensor, Dict]:
-        """NN-on-the-sphere head (T5-chained with the byte tail):
-        in-bank:  -log gate - log softmax_{target ∪ kNN(target)}(T * yhat·c)
-                  + lambda_cos (1 - yhat·c_target)            [aiming term]
-        out-bank: -log(1-gate) - log P_byte(g)
-        Decode metric: exact-NN rate + median angular error vs the budget
-        theta_NN/2 (the decode-correctness theorem)."""
-        cfg = self.cfg
-        logs: Dict[str, float] = {}
-        coords = self._cand_coords()                         # (M, d_point)
-        if step - self._nn_step >= cfg.pointer_refresh or len(self._nn_table) == 0:
-            self._refresh_nn(coords.detach(), step)
-
-        yhat = F.normalize(self.W_point(h), dim=-1)          # (B,S,d_point)
-        bidx = self._bank_index(targets)
-        in_bank = bidx >= 0
-        logs["coverage"] = in_bank.float().mean().item()
-
-        g_logit = self.gate(h).squeeze(-1)
-        nll_byte = -self._byte_logprob(h, targets)
-
-        B, S = bidx.shape
-        tgt = bidx.clamp_min(0)                              # (B,S)
-        negs = self._nn_table[tgt]                           # (B,S,k) hard negatives
-        cand_idx = torch.cat([tgt.unsqueeze(-1), negs], dim=-1)   # (B,S,1+k)
-        c = coords[cand_idx]                                 # (B,S,1+k,d_point)
-        logits = torch.einsum("bsd,bsnd->bsn", yhat, c) * self.point_T
-        nll_point = F.cross_entropy(
-            logits.reshape(-1, logits.shape[-1]),
-            torch.zeros(B * S, dtype=torch.long, device=h.device),
-            reduction="none").view(B, S)
-        cos_t = torch.einsum("bsd,bsd->bs", yhat, coords[tgt])
-        aim = cfg.pointer_cos_weight * (1.0 - cos_t)
-
-        nll = torch.where(in_bank,
-                          -F.logsigmoid(g_logit) + nll_point + aim,
-                          -F.logsigmoid(-g_logit) + nll_byte)
-        loss = nll.mean()
-        logs["bpb"] = loss.item() / 3 / math.log(2)
-        logs["gate_acc"] = ((torch.sigmoid(g_logit) > 0.5) == in_bank
-                            ).float().mean().item()
-        with torch.no_grad():                                # decode metrics
-            if in_bank.any():
-                full = (yhat @ coords.t())                   # (B,S,M)
-                pred = full.argmax(-1)
-                logs["nn_exact"] = (pred[in_bank] == tgt[in_bank]
-                                    ).float().mean().item()
-                # PROPER eval likelihood: full-bank softmax (comparable to
-                # hybrid bpb; the training loss above is contrastive-over-33
-                # and is NOT a likelihood — do not compare it across heads)
-                nll_full = F.cross_entropy(
-                    (full * self.point_T).reshape(-1, full.shape[-1]),
-                    tgt.reshape(-1), reduction="none").view_as(tgt)
-                nll_eval = torch.where(in_bank,
-                                       -F.logsigmoid(g_logit) + nll_full,
-                                       -F.logsigmoid(-g_logit) + nll_byte)
-                logs["bpb_eval"] = nll_eval.mean().item() / 3 / math.log(2)
-                ang = torch.acos(cos_t[in_bank].clamp(-1, 1)) * 180 / math.pi
-                logs["ang_err_deg"] = ang.median().item()
-                logs["budget_deg"] = self._decode_budget_deg
-                logs["in_budget"] = (ang < self._decode_budget_deg
-                                     ).float().mean().item()
-        return loss, logs
-
-    # ---- the loss (T5-exact hybrid + auxiliaries) ----
-    def forward_loss(self, ids: Tensor, targets: Tensor,
-                     step: int = 0) -> Tuple[Tensor, Dict]:
-        cfg = self.cfg
-        h = self.backbone(ids)                               # (B,S,d)
-        logs: Dict[str, float] = {}
-
-        if cfg.head == "pointer":
-            assert self.bank is not None, "pointer head requires a bank"
-            return self._pointer_loss(h, targets, step)
-
-        if cfg.head == "byte":
-            nll = -self._byte_logprob(h, targets)            # (B,S)
-            loss = nll.mean()
-            logs["bpb"] = loss.item() / 3 / math.log(2)
-            return loss, logs
-
-        pi_p, pi_m = self._pi(h)                             # (B,S,K) each
-
-        if cfg.head == "sampled":
-            # uniform negatives + target; uniform proposal => logQ constant,
-            # cancels in softmax (literature requirement satisfied trivially)
-            B, S, _ = h.shape
-            neg = torch.randint(0, 256, (cfg.n_negatives, 3), device=h.device)
-            cand = torch.cat([targets.reshape(-1, 3), neg], dim=0)
-            cand_ids, inv = torch.unique(_tri_ids(cand), return_inverse=True)
-            uniq = torch.stack([cand_ids // 65536, (cand_ids // 256) % 256,
-                                cand_ids % 256], dim=-1)
-            k_p, k_m = self._kappa(uniq)
-            logits = self._bank_logits(pi_p.reshape(-1, cfg.K),
-                                       pi_m.reshape(-1, cfg.K), k_p, k_m)
-            tgt_idx = inv[: B * S]
-            loss = F.cross_entropy(logits, tgt_idx)
-            logs["bpb"] = loss.item() / 3 / math.log(2)
-            logs["n_cand"] = float(len(uniq))
-            return loss, logs
-
-        # banked heads
-        assert self.bank is not None, "head='hybrid'/'bank' requires a bank"
-        if cfg.head == "hybrid" and cfg.bank_scorer == "pmix":
-            logits, pm_logs = self._pmix_logits(h)           # (B,S,M)
-            logs.update(pm_logs)
-        else:
-            k_p, k_m = self._kappa(self.bank)                # (M,K) each
-            logits = self._bank_logits(pi_p, pi_m, k_p, k_m) # (B,S,M)
-        bidx = self._bank_index(targets)                     # (B,S), -1 = miss
-        in_bank = bidx >= 0
-        logs["coverage"] = in_bank.float().mean().item()
-
-        if cfg.head == "bank":
-            # ablation head: proper only on covered targets (coverage logged)
-            lb = F.log_softmax(logits, dim=-1)
-            nll = -lb.gather(-1, bidx.clamp_min(0).unsqueeze(-1)).squeeze(-1)
-            loss = nll[in_bank].mean() if in_bank.any() else logits.sum() * 0
-            logs["bpb_inbank"] = (loss.item() / 3 / math.log(2)
-                                  if in_bank.any() else float("nan"))
-            return loss, logs
-
-        # ── T5-exact hybrid:  -log P(g) per position ──
-        g_logit = self.gate(h).squeeze(-1)                   # (B,S)
-        log_g = F.logsigmoid(g_logit)
-        log_1mg = F.logsigmoid(-g_logit)
-        lb = F.log_softmax(logits, dim=-1)
-        nll_bank = -lb.gather(-1, bidx.clamp_min(0).unsqueeze(-1)).squeeze(-1)
-        nll_byte = -self._byte_logprob(h, targets)
-        nll = torch.where(in_bank, -log_g + nll_bank, -log_1mg + nll_byte)
-        loss = nll.mean()
-        logs["bpb"] = loss.item() / 3 / math.log(2)
-        with torch.no_grad():                       # branch-conditional currencies
-            if in_bank.any():
-                logs["bpb_bank_cond"] = nll_bank[in_bank].mean().item() / 3 / math.log(2)
-            if (~in_bank).any():
-                logs["bpb_byte_cond"] = nll_byte[~in_bank].mean().item() / 3 / math.log(2)
-        logs["gate_acc"] = ((torch.sigmoid(g_logit) > 0.5) == in_bank
-                            ).float().mean().item()
-
-        # ── auxiliaries ──
-        if cfg.write_weight > 0:
-            dz, valid = self._write_target(ids)
-            pred = F.log_softmax(self.W_write(h), dim=-1)
-            kl = F.kl_div(pred, dz, reduction="none").sum(-1)
-            wl = kl[:, valid].mean()
-            loss = loss + cfg.write_weight * wl
-            logs["write_kl"] = wl.item()
-
-        return loss, logs
-
-    # ---- the branching gauge ([TAU] inverted) ----
-    @torch.no_grad()
-    def branching_gauge(self, ids: Tensor, n_baseline: int = 4096) -> Dict:
-        h = self.backbone(ids)
-        pi_p, pi_m = self._pi(h)
-        A = F.normalize(self.codebook, dim=-1)
-        conf = ((pi_p - pi_m) @ A).norm(dim=-1).reshape(-1)
-        rows = F.normalize(torch.randn(n_baseline, self.cfg.D_addr,
-                                       device=h.device), dim=-1)
-        bp, bm = self._address_rows(rows)
-        base = ((bp - bm) @ A).norm(dim=-1)
-        mu, sd = base.mean(), base.std()
-        return {"conf_mean": conf.mean().item(),
-                "kernel_invariant": mu.item(),
-                "branching_frac": (conf < mu - 2 * sd).float().mean().item()}
-
-
-# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-# Training
-# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-
-def train_aleph_lm(cfg: AlephLMConfig,
-                   stream: Optional[TrigramStream] = None) -> Dict:
-    torch.manual_seed(cfg.seed)
-    dev = torch.device(cfg.device)
-    stream = stream or TrigramStream(cfg.corpus_id, cfg.split,
-                                     cfg.max_corpus_bytes, cfg.seed)
-    bank = None
-    if cfg.head in ("hybrid", "bank", "pointer"):
-        if os.path.isfile(str(cfg.bank_source)):           # stratified-atlas bank
-            d = torch.load(cfg.bank_source, map_location="cpu", weights_only=False)
-            bank = d["bank"] if isinstance(d, dict) else d
-            print(f"[bank] loaded {len(bank)} trigram candidates "
-                  f"from {cfg.bank_source}")
-        elif cfg.bank_source == "wordnet":
-            try:
-                bank = build_wordnet_bank(cfg.bank_size)
-                print(f"[bank] wordnet char_eng_3gram: {len(bank)} types")
-            except Exception as e:
-                print(f"[bank] wordnet unavailable ({e}); falling back to corpus")
-        if bank is None:
-            bank = build_corpus_bank(stream, cfg.bank_size)
-            print(f"[bank] corpus top-{len(bank)} trigrams")
-
-    model = AlephLM(cfg, bank=bank).to(dev)
-    n_params = sum(p.numel() for p in model.parameters())
-    opt = torch.optim.Adam(model.parameters(), lr=cfg.lr)    # pure Adam
-    sched = (torch.optim.lr_scheduler.CosineAnnealingLR(
-        opt, T_max=cfg.steps, eta_min=cfg.lr * 0.1) if cfg.lr_decay else None)
-    alephs = model.aleph_layers()
-    for a in alephs:
-        a.emit_diversity = cfg.div_weight > 0
-
-    snapshots: List[Tuple[int, Tensor]] = []
-    if cfg.snapshot_codebook:
-        snapshots.append((0, alephs[0].export_codebook()))
-
-    print(f"\n=== AlephLM  head={cfg.head}  pi={cfg.pi_mode}  "
-          f"bank={cfg.bank_source if bank is not None else '-'}  "
-          f"params={n_params:,}  ctx={cfg.seq_len} tri  "
-          f"eff.batch={cfg.batch_size * cfg.accum_steps}  dev={dev} ===")
-    result: Dict = {"head": cfg.head, "params": n_params}
-    t0 = time.time()
-
-    for step in range(1, cfg.steps + 1):
-        opt.zero_grad(set_to_none=True)
-        loss_sum, logs_acc = 0.0, {}
-        for _ in range(cfg.accum_steps):
-            ids, targets = stream.sample(cfg.batch_size, cfg.seq_len, dev)
-            loss, logs = model.forward_loss(ids, targets, step=step)
-            total = loss
-            if cfg.div_weight > 0:
-                total = total + cfg.div_weight * sum(
-                    a.diversity_loss() for a in alephs)
-            (total / cfg.accum_steps).backward()
-            loss_sum += loss.item()
-            logs_acc = logs
-        loss_avg = loss_sum / cfg.accum_steps
-        gnorm = torch.nn.utils.clip_grad_norm_(
-            model.parameters(), max(loss_avg, 1.0))
-        opt.step()
-        if sched is not None:
-            sched.step()
-
-        if step % cfg.log_every == 0 or step == cfg.steps:
-            rate = step * cfg.batch_size * cfg.seq_len * cfg.accum_steps \
-                / (time.time() - t0)
-            line = (f"  step {step:6d}  loss {loss_avg:.4f}  "
-                    f"bpb {logs_acc.get('bpb', logs_acc.get('bpb_inbank', float('nan'))):.3f}  "
-                    f"|g| {gnorm:.2f}  {rate/1e3:.1f}k tri/s")
-            if "coverage" in logs_acc:
-                line += f"  cov {logs_acc['coverage']:.0%}"
-            if "gate_acc" in logs_acc:
-                line += f"  gate {logs_acc['gate_acc']:.0%}"
-            if "nn_exact" in logs_acc:
-                line += (f"  bpbE {logs_acc.get('bpb_eval', float('nan')):.3f}"
-                         f"  nn {logs_acc['nn_exact']:.0%}"
-                         f"  ang {logs_acc['ang_err_deg']:.1f}/"
-                         f"{logs_acc['budget_deg']:.1f}deg"
-                         f"  inBudget {logs_acc['in_budget']:.0%}")
-            if "bpb_bank_cond" in logs_acc:
-                line += (f"  inB {logs_acc['bpb_bank_cond']:.3f}"
-                         f"  outB {logs_acc.get('bpb_byte_cond', float('nan')):.3f}")
-            if "mode_spread_deg" in logs_acc:
-                line += (f"  spread {logs_acc['mode_spread_deg']:.0f}deg"
-                         f"  mixH {logs_acc['mix_entropy']:.2f}")
-            if "write_kl" in logs_acc:
-                line += f"  wKL {logs_acc['write_kl']:.3f}"
-            model.eval()
-            with torch.no_grad():
-                ids_p, _ = stream.sample(min(8, cfg.batch_size), cfg.seq_len, dev)
-                st = alephs[0].address_stats(model.backbone(ids_p),
-                                             max_rows=200_000)
-                bg = model.branching_gauge(ids_p)
-            model.train()
-            line += (f"  ppl {st['perplexity']:.0f}/{st['max_perplexity']:.0f}"
-                     f"  conf {bg['conf_mean']:.3f}"
-                     f"/{bg['kernel_invariant']:.3f}"
-                     f"  branch {bg['branching_frac']:.0%}")
-            print(line)
-            result.update(logs_acc)
-            result.update({"loss": loss_avg, "step": step, **bg})
-            if cfg.snapshot_codebook:
-                snapshots.append((step, alephs[0].export_codebook()))
-
-    if snapshots:
-        traj = [(s, statute(cb)) for s, cb in snapshots]
-        result["statute_trajectory"] = traj
-        torch.save({"snapshots": snapshots, "statute_trajectory": traj,
-                    "config": cfg.__dict__}, cfg.snapshot_path)
-        d0, d1 = traj[0][1]["deviation"], traj[-1][1]["deviation"]
-        print(f"\n[basin] statute: dev {d0:+.4f} -> {d1:+.4f} "
-              f"({traj[-1][1]['statute']});  snapshots -> {cfg.snapshot_path}")
-    if cfg.checkpoint_path:
-        torch.save({"model_state_dict": model.state_dict(),
-                    "config": cfg.__dict__,
-                    "bank": model.bank.cpu() if model.bank is not None else None},
-                   cfg.checkpoint_path)
-        print(f"[ckpt] -> {cfg.checkpoint_path}")
-    return result
-
-
-# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-# Smoke + activation
-# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-
-def _smoke():
-    print("=" * 70)
-    print("AlephLM — smoke")
-    print("=" * 70)
-    rng = np.random.default_rng(0)
-    words = [b"the", b"aleph", b"predicts", b"its", b"own", b"future",
-             b"through", b"a", b"codebook"]
-    path = "/tmp/_alm_corpus.txt"
-    with open(path, "wb") as f:
-        f.write(b" ".join(words[i] for i in rng.integers(0, 9, 80000)))
-
-    base = dict(corpus_id=path, max_corpus_bytes=None, steps=25, log_every=25,
-                dim=96, n_layers=2, n_heads=4, K=16, seq_len=48, batch_size=8,
-                bank_size=256, n_negatives=128, device="cpu",
-                checkpoint_path=None, snapshot_path="/tmp/_alm_snaps.pt")
-    prior_bpb = 8.0
-    for head in ("hybrid", "byte", "bank", "sampled"):
-        r = train_aleph_lm(AlephLMConfig(head=head, **base))
-        bpb = r.get("bpb", r.get("bpb_inbank", float("nan")))
-        assert math.isfinite(r["loss"]), head
-        print(f"  ✓ head={head:8s} loss {r['loss']:.3f}  bpb {bpb:.2f} "
-              f"(uniform prior {prior_bpb:.1f})")
-    # pi ablation path + gradient to codebook through the PREDICT/CANDIDATE legs
-    cfg = AlephLMConfig(head="hybrid", pi_mode="address", **base)
-    stream = TrigramStream(path, max_corpus_bytes=None, seed=0)
-    bank = build_corpus_bank(stream, cfg.bank_size)
-    m = AlephLM(cfg, bank=bank)
-    ids, tg = stream.sample(4, cfg.seq_len, "cpu")
-    loss, _ = m.forward_loss(ids, tg)
-    loss.backward()
-    assert m.codebook.grad is not None and torch.isfinite(m.codebook.grad).all()
-    print(f"  ✓ pi_mode='address' ablation runs; codebook grad |{m.codebook.grad.norm():.3f}|")
-    print("All smoke tests passed.")
-
-
-#if __name__ == "__main__":
-#    import argparse
-#    ap = argparse.ArgumentParser(description="AlephLM — prediction through the codebook")
-#    ap.add_argument("--smoke-only", action="store_true")
-#    ap.add_argument("--head", default="hybrid",
-#                    choices=["hybrid", "byte", "bank", "sampled", "pointer"])
-#    ap.add_argument("--bank", default="corpus", choices=["corpus", "wordnet"])
-#    ap.add_argument("--pi", default="free", choices=["free", "address"])
-#    ap.add_argument("--scorer", default="kernel", choices=["kernel", "pmix"])
-#    ap.add_argument("--pointers", type=int, default=4)
-#    ap.add_argument("--steps", type=int, default=10_000)
-#    ap.add_argument("--corpus-mb", type=int, default=100)
-#    ap.add_argument("--device",
-#                    default="cuda" if torch.cuda.is_available() else "cpu")
-#    args, _unknown = ap.parse_known_args()
-#    if args.smoke_only:
-#        _smoke()
-#    else:
-#        cfg = AlephLMConfig(head=args.head, bank_source=args.bank,
-#                            pi_mode=args.pi, steps=args.steps,
-#                            bank_scorer=args.scorer, n_pointers=args.pointers,
-#                            max_corpus_bytes=args.corpus_mb * 1_000_000,
-#                            device=args.device)
-#        train_aleph_lm(cfg)