Update experiments/exp_007_aleph_routed_attention/4_aleph_lm.py

experiments/exp_007_aleph_routed_attention/4_aleph_lm.py

@@ -0,0 +1,1031 @@
+# 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
+
+from aleph_routed_attention import AlephRoutedAttention, AlephAttentionConfig
+from aleph_trigram_lm import TrigramStream, statute
+
+
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+# 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'
+    # ── Tier-A scaling switches (2026-06-11) ──
+    # bank_softmax='sampled' (pmix only): train CE over {batch targets} ∪
+    # n_bank_samples uniform negatives with the log-Q correction (Jean et al.;
+    # shared-negative approximation documented at the call site). The REPORTED
+    # bpb stays honest: at every log step the in-bank NLL is recomputed against
+    # the FULL bank in chunked no_grad. pos_mode='clamp' saturates position
+    # indices at seq_len-1 -> unbounded streaming/generation (decay-gated state
+    # is the principled v2). train_mode='stream' = TBPTT-1 over `segments`
+    # carried streaming states: effective context segments*seq_len at constant
+    # memory (requires pos_mode='clamp').
+    bank_softmax: str = "full"               # 'full' | 'sampled'
+    n_bank_samples: int = 8192
+    pos_mode: str = "absolute"               # 'absolute' | 'clamp'
+    train_mode: str = "window"               # 'window' | 'stream'
+    segments: int = 4
+    amp: bool = True                         # bf16 autocast on CUDA
+    compile_backbone: bool = False           # compile backbone only (tensor out)
+    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_softmax in ("full", "sampled")
+        assert self.pos_mode in ("absolute", "clamp")
+        assert self.train_mode in ("window", "stream")
+        if self.bank_softmax == "sampled":
+            assert self.head == "hybrid" and self.bank_scorer == "pmix", \
+                "sampled softmax is wired for the pmix hybrid"
+        if self.train_mode == "stream":
+            assert self.pos_mode == "clamp", "stream training needs pos_mode='clamp'"
+        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,
+                     cols: Optional[Tensor] = None) -> Tuple[Tensor, Dict[str, float]]:
+        cfg = self.cfg
+        J = cfg.n_pointers
+        coords = self._cand_coords()                          # (M, d_point)
+        if cols is not None:
+            coords = coords[cols]                             # (Mc, 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}
+
+    # ---- position (absolute, or clamped for unbounded streaming) ----
+    def _pos_slice(self, S: int, offset: int = 0) -> Tensor:
+        if self.cfg.pos_mode == "clamp":
+            idx = (torch.arange(S, device=self.pos.device) + offset
+                   ).clamp_max(self.cfg.seq_len - 1)
+            return self.pos[:, idx]
+        return self.pos[:, offset: offset + S]
+
+    # ---- honest full-bank in-bank NLL (chunked, no_grad) for sampled mode ----
+    @torch.no_grad()
+    def _pmix_full_nll(self, h: Tensor, bidx: Tensor, in_bank: Tensor) -> Tensor:
+        M = self.bank.shape[0]
+        lse = None
+        tgt_logit = torch.zeros_like(bidx, dtype=h.dtype)
+        for lo in range(0, M, 8192):
+            cols = torch.arange(lo, min(lo + 8192, M), device=h.device)
+            lg, _ = self._pmix_logits(h, cols=cols)           # (B,S,Mc)
+            chunk_lse = torch.logsumexp(lg, dim=-1)
+            lse = chunk_lse if lse is None else torch.logaddexp(lse, chunk_lse)
+            hit = in_bank & (bidx >= lo) & (bidx < lo + cols.numel())
+            if hit.any():
+                tgt_logit[hit] = lg[hit].gather(
+                    -1, (bidx[hit] - lo).unsqueeze(-1)).squeeze(-1)
+        return lse - tgt_logit                                # NLL (B,S)
+
+    # ---- tower ----
+    def backbone(self, ids: Tensor) -> Tensor:
+        x = sum(emb(ids[..., i]) for i, emb in enumerate(self.byte_emb))
+        x = x + self._pos_slice(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, h: Optional[Tensor] = None
+                     ) -> Tuple[Tensor, Dict]:
+        cfg = self.cfg
+        if h is None:
+            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"
+        bidx = self._bank_index(targets)                     # (B,S), -1 = miss
+        in_bank = bidx >= 0
+        logs["coverage"] = in_bank.float().mean().item()
+        sampled = (cfg.bank_softmax == "sampled" and self.training)
+        if cfg.head == "hybrid" and cfg.bank_scorer == "pmix":
+            if sampled:
+                # shared-negative sampled softmax: cols = batch targets ∪
+                # uniform negatives; log-Q correction log(n/M) on pure
+                # negatives only (cols that are some example's target carry
+                # Q≈1; tiny documented bias — reported bpb is recomputed
+                # full-bank below at log steps, so the LOGGED number is exact)
+                M = self.bank.shape[0]
+                n = min(cfg.n_bank_samples, M)
+                tcols = bidx[in_bank].unique()
+                samp = torch.randint(0, M, (n,), device=h.device)
+                cols = torch.unique(torch.cat([tcols, samp]))
+                logits, pm_logs = self._pmix_logits(h, cols=cols)
+                logits = logits + torch.where(
+                    torch.isin(cols, tcols), 0.0,
+                    -math.log(n / M)).to(logits.dtype)        # -(-logQ)
+                bidx = torch.searchsorted(cols, bidx.clamp_min(0))
+            else:
+                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)
+
+        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()
+        if sampled and (step % max(cfg.log_every, 1) == 0):
+            nb_full = self._pmix_full_nll(h.detach(), self._bank_index(targets),
+                                          in_bank)
+            nll_h = torch.where(in_bank, -log_g.detach() + nb_full,
+                                -log_1mg.detach() + nll_byte.detach())
+            logs["bpb"] = nll_h.mean().item() / 3 / math.log(2)   # honest
+            logs["bpb_s"] = loss.item() / 3 / math.log(2)         # sampled
+        else:
+            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
+
+
+    # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+    # Branch-predictive generation (streaming, gauge-triggered forking)
+    # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+    # The CPU metaphor, implemented: per-branch state is the streaming
+    # (M±, z±) per layer — O(K·d) regardless of context, so speculation is
+    # cheap. The model's own multimodality estimate is the branch-predictor
+    # confidence bit: fork ONLY where the next-distribution is genuinely
+    # multimodal (top-2 probability ratio above `fork_ratio`); run straight
+    # through deterministic stretches at full pipeline speed.
+    # Position embedding is absolute: generation is capped at cfg.seq_len
+    # total positions (documented v1 limitation).
+
+    def stream_backbone(self, ids: Tensor, states: Optional[List],
+                        pos_offset: int) -> Tuple[Tensor, List]:
+        """One streamed segment (B, S_seg, 3) with per-layer carried states."""
+        x = sum(emb(ids[..., i]) for i, emb in enumerate(self.byte_emb))
+        x = x + self._pos_slice(ids.shape[1], pos_offset)
+        states = states or [None] * len(self.layers)
+        new_states: List = []
+        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)
+        return self.norm_f(x), new_states
+
+    @torch.no_grad()
+    def next_distribution(self, h_last: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+        """(gate_prob, bank_probs (M,), byte_logprobs (3,256)) for one position.
+        h_last: (1, d). Bank scorer follows cfg.bank_scorer."""
+        cfg = self.cfg
+        g = torch.sigmoid(self.gate(h_last)).squeeze()
+        if cfg.bank_scorer == "pmix":
+            logits, _ = self._pmix_logits(h_last.unsqueeze(0))
+            bank_p = F.softmax(logits.squeeze(0).squeeze(0), dim=-1)
+        else:
+            pi_p, pi_m = self._pi(h_last)
+            k_p, k_m = self._kappa(self.bank)
+            bank_p = F.softmax(self._bank_logits(pi_p, pi_m, k_p, k_m
+                                                 ).squeeze(0), dim=-1)
+        byte_lp = torch.stack([F.log_softmax(head(h_last).squeeze(0), dim=-1)
+                               for head in self.byte_heads])         # (3,256)
+        return g, bank_p, byte_lp
+
+    @torch.no_grad()
+    def generate_tree(self, prompt: bytes, max_new: int = 24,
+                      beam: int = 8, fork_ratio: float = 0.35,
+                      fork_width: int = 3, device: str = "cpu") -> List[Dict]:
+        """Branch-predictive decoding. Returns the surviving branches as
+        [{'text', 'logp', 'forks'}], best first.
+        fork_ratio: fork iff p2/p1 > ratio (the confidence bit);
+        fork_width: children per fork; beam: global survivor cap."""
+        self.eval()
+        cfg = self.cfg
+        b = prompt[: 3 * (len(prompt) // 3)] or b"   "
+        ids = torch.frombuffer(bytearray(b), dtype=torch.uint8) \
+                   .to(torch.long).view(1, -1, 3).to(device)
+        h, states = self.stream_backbone(ids, None, 0)
+        pos = ids.shape[1]
+        Branch = lambda st, h_, lp, txt, forks: \
+            {"states": st, "h": h_, "logp": lp, "bytes": txt, "forks": forks}
+        branches = [Branch(states, h[:, -1], 0.0, b"", 0)]
+
+        bank_bytes = self.bank.cpu().numpy().astype("uint8") \
+            if self.bank is not None else None
+        for step in range(max_new):
+            if cfg.pos_mode == "absolute" and pos + 1 > cfg.seq_len:
+                break
+            nxt: List[Dict] = []
+            for br in branches:
+                g, bank_p, byte_lp = self.next_distribution(br["h"])
+                # mixture distribution over candidate continuations:
+                # in-bank candidates carry g*bank_p; the byte tail is
+                # summarized by its argmax trigram carrying (1-g)*p_byte
+                cand_p, cand_tri = [], []
+                if bank_bytes is not None:
+                    top_p, top_i = bank_p.topk(min(fork_width + 1, len(bank_p)))
+                    for p, i in zip(top_p.tolist(), top_i.tolist()):
+                        cand_p.append(g.item() * p)
+                        cand_tri.append(bytes(bank_bytes[i]))
+                by = byte_lp.argmax(-1)
+                p_by = float(byte_lp.max(-1).values.sum().exp())
+                cand_p.append((1 - g.item()) * p_by)
+                cand_tri.append(bytes(by.tolist()))
+                order = np.argsort(cand_p)[::-1]
+                p1 = cand_p[order[0]]
+                p2 = cand_p[order[1]] if len(order) > 1 else 0.0
+                take = order[: fork_width] if (p1 > 0 and p2 / max(p1, 1e-12)
+                                               > fork_ratio) else order[:1]
+                forked = len(take) > 1
+                for oi in take:
+                    tri = cand_tri[oi]
+                    t = torch.tensor(list(tri), dtype=torch.long,
+                                     device=device).view(1, 1, 3)
+                    h2, st2 = self.stream_backbone(
+                        t, [tuple(s.clone() for s in st)
+                            for st in br["states"]], pos)
+                    nxt.append(Branch(st2, h2[:, -1],
+                                      br["logp"] + math.log(max(cand_p[oi], 1e-12)),
+                                      br["bytes"] + tri,
+                                      br["forks"] + int(forked)))
+            nxt.sort(key=lambda d: -d["logp"])
+            branches = nxt[:beam]
+            pos += 1
+        return [{"text": (prompt + br["bytes"]).decode("utf-8", errors="replace"),
+                 "logp": br["logp"], "forks": br["forks"]}
+                for br in branches]
+
+    # ---- 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()
+
+    use_amp = cfg.amp and dev.startswith("cuda")
+    if cfg.compile_backbone:
+        model.backbone = torch.compile(model.backbone)   # tensor-out only
+
+    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):
+            if cfg.train_mode == "stream":
+                ids, targets = stream.sample(
+                    cfg.batch_size, cfg.seq_len * cfg.segments, dev)
+                states = None
+                for si in range(cfg.segments):
+                    sl = slice(si * cfg.seq_len, (si + 1) * cfg.seq_len)
+                    with torch.autocast(device_type="cuda",
+                                        dtype=torch.bfloat16, enabled=use_amp):
+                        hseg, states = model.stream_backbone(
+                            ids[:, sl], states, si * cfg.seq_len)
+                        loss, logs = model.forward_loss(
+                            ids[:, sl], targets[:, sl], step=step, h=hseg)
+                        total = loss
+                        if cfg.div_weight > 0:
+                            total = total + cfg.div_weight * sum(
+                                a.diversity_loss() for a in alephs)
+                    (total / cfg.accum_steps / cfg.segments).backward()
+                    states = [tuple(t.detach() for t in st) for st in states]
+                    loss_sum += loss.item() / cfg.segments
+                logs_acc = logs
+                continue
+            ids, targets = stream.sample(cfg.batch_size, cfg.seq_len, dev)
+            with torch.autocast(device_type="cuda",
+                                dtype=torch.bfloat16, enabled=use_amp):
+                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 "bpb_s" in logs_acc:
+                line += f"  bpbS {logs_acc['bpb_s']:.3f}"
+            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.")
+
+
+def tier_a_config(bank_source: str, **overrides) -> AlephLMConfig:
+    """Tier A of the scaling plan (~25-30M params, one Blackwell, days):
+    d=512 x 8 layers, J=16 pmix on a dense bank with sampled softmax,
+    TBPTT streaming (4 x 1024 = effective 4096-trigram context at constant
+    memory), clamped position (unbounded generation), bf16. Pure Adam,
+    standing clip rule — nothing exotic enters with scale."""
+    base = dict(dim=512, n_layers=8, n_heads=8, K=64, d_point=48,
+                head="hybrid", bank_scorer="pmix", n_pointers=16,
+                bank_source=bank_source, bank_softmax="sampled",
+                n_bank_samples=8192, pos_mode="clamp", train_mode="stream",
+                segments=4, seq_len=1024, batch_size=16, accum_steps=4,
+                lr=5e-4, steps=50_000, amp=True, log_every=250)
+    base.update(overrides)
+    return AlephLMConfig(**base)
+
+
+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)