Update experiments/exp_007_aleph_routed_attention/aleph_routed_trigram_lm.py

experiments/exp_007_aleph_routed_attention/aleph_routed_trigram_lm.py

@@ -1,60 +1,67 @@
-# aleph_trigram_lm.py
+# aleph_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
+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
@@ -62,7 +69,7 @@ from __future__ import annotations
 import math
 import os
 import time
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Dict, List, Optional, Tuple
 
 import numpy as np
@@ -76,250 +83,292 @@ from torch import Tensor
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 
 @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.
-    """
+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 (= 768 bytes of context)
+    seq_len: int = 256                       # trigrams (3*seq_len bytes)
     seed: int = 1234
 
-    # model
+    # tower
     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)
+    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 = 2600                     # optimizer steps (micro-batches =
-                                            # steps * accum_steps)
-    batch_size: int = 512
+    steps: int = 10_000
+    batch_size: int = 32
+    accum_steps: int = 1
     lr: float = 3e-4
     lr_decay: bool = True
-    log_every: int = 100
-    eval_batches: int = 8
+    div_weight: float = 0.0
+    log_every: int = 250
     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"
+    snapshot_path: str = "aleph_lm_snaps.pt"
+    checkpoint_path: Optional[str] = "aleph_lm.pt"
 
     def __post_init__(self):
-        assert self.attn_mode in ("hub", "bucket", "standard")
+        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.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)"
-
-
+        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"
 
 
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-# Statute monitor — the program's own diagnostic geometry (Sec 3.11)
+# Candidate banks
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 
-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}
+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)
 
 
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-# Substrate — trigram stream (mirrors ByteTrigramDataset's loading)
+# Model
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 
-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)
-
+class AlephLM(nn.Module):
+    """The composite reduction. forward_loss(ids, targets) -> (loss, logs)."""
 
-# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-# 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):
+    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))
 
-        def make_attn() -> nn.Module:
-            if cfg.attn_mode == "standard":
-                return StandardAttention(d, cfg.n_heads, causal=True)
+        # ── MIX (hub tower, shared codebook) ──
+        def make_attn():
             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,
+                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)
-        ])
-        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"):
+            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:
-        """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:
@@ -327,206 +376,417 @@ class TrigramLM(nn.Module):
             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)."""
+    # ---- 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]]
-        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
+        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_trigram_lm(cfg: TrigramLMConfig,
-                     stream: Optional[TrigramStream] = None) -> Dict:
+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)
-    model = TrigramLM(cfg).to(dev)
+    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, never AdamW
+    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 and alephs:
+    if cfg.snapshot_codebook:
         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}
+    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()
 
-    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
+        loss_sum, logs_acc = 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
+            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))           # the clip rule
+            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:
-            bpb = loss_avg / math.log(2)
-            rate = (step * cfg.batch_size * cfg.seq_len * segs
-                    * cfg.accum_steps) / (time.time() - t0)
+            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 {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()))
+                    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({"loss": loss_avg, "bpb": bpb, "step": step})
+            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:
-        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).")
+                    "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__}, cfg.checkpoint_path)
-        print(f"[ckpt] saved -> {cfg.checkpoint_path}")
+                    "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
 
 
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-# Activation
+# Smoke + activation
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 
-def _smoke(device: str = "cpu"):
-    """End-to-end on a synthetic local corpus — no downloads."""
+def _smoke():
     print("=" * 70)
-    print("aleph_trigram_lm — smoke (synthetic corpus)")
+    print("AlephLM — smoke")
     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"]
+    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, 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")
+        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="Aleph trigram LM — basin run")
+#    ap = argparse.ArgumentParser(description="AlephLM — prediction through the codebook")
 #    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=2600)
+#    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("--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)
-#
+#    args, _unknown = ap.parse_known_args()
 #    if args.smoke_only:
-#        _smoke(device="cpu")
+#        _smoke()
 #    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)
+#        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)