Spaces:

Ex0bit
/

mLattice

Running

μLattice: live demo — your own facts, store at scale

20a3a74 verified 28 days ago

22.3 kB

"""Shared evaluation/training protocol for the capacity, forgetting and real-data goals.
Single source of truth: the acceptance tests (tests/test_capacity.py,
tests/test_forgetting.py, tests/test_realtext.py) and the training scripts
(experiments/train_capacity.py, experiments/train_forgetting.py,
experiments/realdata_v2.py) all build sequences and measure accuracy through
this module, so "train until the test passes" cannot drift from what is tested.
Synthetic task (capacity + forgetting):
    vocab  keys [0, NK)  values [NK, NK+NV)  SEP = NK+NV+1   V = NK+NV+2
    sequence per row (filler f tokens before each pair; f=0 -> dense MQAR):
        [f k1 v1] [f k2 v2] ... [f kP vP]            store pass
        [f ku1 vu1'] ... [f kuM vuM']                update pass (overwrites)
        [q1 SEP] [q2 SEP] ... [qQ SEP]               query pass
    A query on an updated key must return the UPDATED value (last write wins).
    The answer is read at the SEP token of each query (causal).
Real-data task: wikitext-103 mined key->bigram-successor pairs embedded in a
real corpus filler stream (same mining as experiments/eval_realdata.py).
"""
from __future__ import annotations

import os
import time
from dataclasses import dataclass

import torch

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.join(HERE, "..")

# ---------------------------------------------------------------------------
# synthetic generator
# ---------------------------------------------------------------------------


@dataclass
class SynthCfg:
    NK: int = 16384          # distinct key tokens available
    NV: int = 16384          # distinct value tokens available
    d_model: int = 64
    n_heads: int = 2
    d_k: int = 32
    d_v: int = 32
    addr_scale: float = 64.0  # (2*64+3)^2 = 17,161 cells for 8,192 keys
    retention_init: float = 9.0

    @property
    def V(self) -> int:
        return self.NK + self.NV + 2

    @property
    def SEP(self) -> int:
        return self.NK + self.NV + 1


def synth_batch(
    B: int,
    n_pairs: int,
    n_queries: int,
    cfg: SynthCfg,
    filler: int = 0,
    overwrite_frac: float = 0.0,
    generator: torch.Generator | None = None,
    qpos: torch.Tensor | None = None,
    updated_idx: torch.Tensor | None = None,
    kw: dict | None = None,
):
    """Build a capacity/forgetting batch. See module docstring for the layout.
    Returns dict with:
      tokens      (B, L) long          the sequence
      ans_pos     (B, Q) long          SEP positions; answer read there
      gold        (B, Q) long          value token id that must be produced
      was_updated (B, Q) bool          query key was overwritten -> updated gold
      key_pos     (B, P) long          first-pass key token positions (spread loss)
      upd_key_pos (B, M) long | None   update-pass key token positions
      n_pairs, n_upd, L
    """
    g = generator
    NK, NV = cfg.NK, cfg.NV
    if n_pairs > NK:
        raise ValueError("n_pairs > key vocab")
    n_upd = int(round(n_pairs * overwrite_frac))
    if n_pairs + n_upd > NV:
        raise ValueError("not enough distinct values")
    dev = "cpu"  # build on CPU; caller moves to device

    # distinct keys per row; `include_keys` are guaranteed members of the stored
    # set (placed FIRST), and the rest are sampled from the complement, so the
    # batch is consistent by construction -- no duplicate stores of the forced
    # keys with conflicting values (the token-surgery approach produced exactly
    # that and wrecked the model: census 9 -> 713 errors, polish_weakkeys v1).
    include_keys = None
    if "include_keys" in (kw or {}):
        include_keys = kw["include_keys"]
        if include_keys.dim() == 1:
            include_keys = include_keys.unsqueeze(0).expand(B, -1)
    if include_keys is not None:
        nk_inc = include_keys.shape[1]
        assert n_pairs >= nk_inc
        rows = []
        for r in range(B):
            rest = torch.randperm(NK, generator=g)
            rest = rest[torch.isin(rest, include_keys[r], invert=True)][: n_pairs - nk_inc]
            rows.append(torch.cat([include_keys[r], rest]))
        keys = torch.stack(rows)
    else:
        keys = torch.stack([torch.randperm(NK, generator=g)[:n_pairs] for _ in range(B)])
    # distinct values: n_pairs originals + n_upd updates, all different in a row
    vals = torch.stack([torch.randperm(NV, generator=g)[: n_pairs + n_upd] for _ in range(B)])
    orig_vals, upd_vals = vals[:, :n_pairs], vals[:, n_pairs:]

    if updated_idx is None:
        # which stored keys get overwritten (row-specific)
        updated_idx = torch.stack([torch.randperm(n_pairs, generator=g)[:n_upd] for _ in range(B)])
    if qpos is None:
        # DISTINCT query keys (sample without replacement).  Sampling with
        # replacement creates duplicate queries at large Q (measured: 9 of 512
        # at Q=512/NK=16,384) -- not part of the published protocol semantics
        # ("every query has a unique answer", mqar.py); distinct sampling is
        # the clean generalization of the published small-Q generator.
        qpos = torch.stack([torch.randperm(n_pairs, generator=g)[:n_queries] for _ in range(B)])
    else:
        if qpos.dim() == 1:
            qpos = qpos.unsqueeze(0).expand(B, -1).contiguous()
        n_queries = qpos.shape[1]      # explicit query set defines Q
    qkeys = torch.gather(keys, 1, qpos)
    upd_mask = torch.zeros(B, n_pairs, dtype=torch.bool)
    upd_mask.scatter_(1, updated_idx, True)
    was_updated = torch.gather(upd_mask, 1, qpos)
    # latest value for every stored key: original, replaced where updated
    latest = orig_vals.clone()
    latest.scatter_(1, updated_idx, upd_vals)
    gold = torch.gather(latest, 1, qpos) + NK

    f = filler
    block = f + 2
    L = (n_pairs + n_upd) * block + 2 * n_queries

    tokens = torch.zeros(B, L, dtype=torch.long)
    # filler ids: uniform from the complement of the row's stored keys
    if f > 0:
        need = f * (n_pairs + n_upd)
        filler_ids = torch.zeros(B, need, dtype=torch.long)
        keyset = torch.zeros(B, NK, dtype=torch.bool)
        keyset.scatter_(1, keys, True)
        done = torch.zeros(B, need, dtype=torch.bool)
        while not bool(done.all()):
            cand = torch.randint(0, NK, (B, need), generator=g)
            ok = ~torch.gather(keyset, 1, cand) & ~done
            filler_ids = torch.where(ok, cand, filler_ids)
            done = done | ok
    # store pass
    for i in range(n_pairs):
        s = i * block
        if f:
            tokens[:, s:s + f] = filler_ids[:, i * f:(i + 1) * f]
        tokens[:, s + f] = keys[:, i]
        tokens[:, s + f + 1] = orig_vals[:, i] + NK
    # update pass: same key, new value, later in the sequence
    upd_tokens = torch.gather(keys, 1, updated_idx)
    for j in range(n_upd):
        s = (n_pairs + j) * block
        if f:
            tokens[:, s:s + f] = filler_ids[:, (n_pairs + j) * f:(n_pairs + j + 1) * f]
        tokens[:, s + f] = upd_tokens[:, j]
        tokens[:, s + f + 1] = upd_vals[:, j] + NK
    # query pass
    S2 = (n_pairs + n_upd) * block
    for q in range(n_queries):
        tokens[:, S2 + 2 * q] = qkeys[:, q]
        tokens[:, S2 + 2 * q + 1] = cfg.SEP

    key_pos = (torch.arange(n_pairs) * block + f).unsqueeze(0).expand(B, -1).contiguous()
    upd_key_pos = ((n_pairs + torch.arange(n_upd)) * block + f).unsqueeze(0).expand(B, -1).contiguous() \
        if n_upd else None
    ans_pos = (S2 + 2 * torch.arange(n_queries) + 1).unsqueeze(0).expand(B, -1).contiguous()
    return {
        "tokens": tokens, "ans_pos": ans_pos, "gold": gold, "was_updated": was_updated,
        "key_pos": key_pos, "upd_key_pos": upd_key_pos, "n_pairs": n_pairs,
        "n_upd": n_upd, "L": L,
    }


# ---------------------------------------------------------------------------
# model helpers
# ---------------------------------------------------------------------------


def build_model(cfg: SynthCfg, n_layers: int = 1):
    from mulattice.models import MQARModel, MixerCfg
    mcfg = MixerCfg(d_model=cfg.d_model, n_heads=cfg.n_heads, d_k=cfg.d_k, d_v=cfg.d_v,
                    addr_scale=cfg.addr_scale, learn_addr_scale=False, bound_address=True,
                    read_tie=False, retention_init=cfg.retention_init)
    return MQARModel(cfg.V, kind="mulattice", n_layers=n_layers, cfg=mcfg)


def logits_at(model, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
    """Forward the trunk, then apply the head ONLY at `positions` (B,Q).
    Avoids materialising (B,T,V) logits; at T=131K with V=32,770 that is 17GB.
    """
    x = model.emb(tokens)
    for blk in model.blocks:
        x = blk(x, read_mode="all")
    x = model.norm(x)
    B = x.shape[0]
    h = x[torch.arange(B, device=x.device).unsqueeze(1), positions]
    return model.head(h)                       # (B,Q,V)


@torch.no_grad()
def eval_batch(model, batch, device) -> dict:
    """Accuracy of `model` on a synth_batch, split by updated/untouched queries."""
    model.eval()
    pred = logits_at(model, batch["tokens"].to(device), batch["ans_pos"].to(device)).argmax(-1).cpu()
    gold = batch["gold"]
    ok = pred == gold
    res = {"acc": ok.float().mean().item(), "n": ok.numel(),
           "correct": int(ok.sum())}
    up = batch["was_updated"]
    if bool(up.any()):
        res["acc_updated"] = ok[up].float().mean().item()
    if bool((~up).any()):
        res["acc_untouched"] = ok[~up].float().mean().item()
    model.train()
    return res


def scan_info(model) -> str:
    parts = []
    for i, blk in enumerate(model.blocks):
        sc = blk.mixer.scan
        parts.append(f"L{i} cells={sc.last_n_cells} rounds={getattr(sc, 'last_rounds', '?')}"
                     f" path={getattr(sc, 'last_path', '?')}")
    return " ".join(parts)


def force_fast_scan(model):
    for blk in model.blocks:
        blk.mixer.scan.crossover = 1e9


def key_addresses(model, tokens, key_pos):
    """Write addresses (B,N,2) of the tokens at `key_pos` (B,N).
    Uses the mixer's LIVE addr_scale (it can be ramped mid-run), not the
    construction-time cfg value."""
    mx = model.blocks[0].mixer
    xc = mx.conv(model.blocks[0].n1(model.emb(tokens)))
    kaddr = mx.cfg.addr_scale * torch.tanh(mx.Wa(xc))
    ar = torch.arange(tokens.shape[0], device=tokens.device).unsqueeze(1)
    return kaddr[ar, key_pos]


def spread_loss(model, tokens, key_pos, cfg: SynthCfg, gate_subsample: int = 2048):
    """Kernel-D gated spread loss on the TRUE key-token write addresses.
    Same definition as mulattice.addressing.gated_spread_loss, but the
    pairwise write-overlap gate is computed on a subsample when N is large
    (it is (B,N,N); at N=8,192 that is 67M entries per row)."""
    addr = key_addresses(model, tokens, key_pos)
    B, N, _ = addr.shape
    addr_gate = addr[:, torch.randperm(N, device=addr.device)[:gate_subsample]] if N > gate_subsample else addr
    return _gated_spread(addr, addr_gate, float(model.blocks[0].mixer.cfg.addr_scale))


def _gated_spread(addr, addr_gate, scale):
    from mulattice.addressing import write_overlap
    from mulattice.kernel import phi_2d
    B, N, _ = addr.shape
    gate = (write_overlap(addr_gate) - 0.05).clamp(0.0, 1.0).detach()
    lam = 0.3 + 0.7 * gate
    w, z = phi_2d(addr)
    size = int(2 * scale + 3)
    off = int(scale) + 1
    gi = (z + off).clamp(0, size - 1)
    idx = (gi[..., 0] * size + gi[..., 1]).reshape(B, -1).to(torch.int64)
    mass = (w / w.sum(dim=(-2, -1), keepdim=True)).reshape(B, -1)
    hist = torch.zeros(B, size * size, device=addr.device, dtype=addr.dtype)
    for b in range(B):
        hist[b].index_add_(0, idx[b], mass[b])
    p = hist / hist.sum(dim=1, keepdim=True).clamp_min(1e-9)
    entropy = -(p * (p + 1e-12).log()).sum(dim=1)
    return (lam * (-entropy)).mean()


# ---------------------------------------------------------------------------
# vocab surgery: extend the proven 1024-key model to the full key vocabulary
# ---------------------------------------------------------------------------


def extend_from_checkpoint(new_model, ckpt_path: str, NK: int = 16384, NV: int = 16384,
                           old_NK: int = 1024, old_NV: int = 1024,
                           noise: float = 0.05, seed: int = 0):
    """Copy the proven model's mixer into `new_model`; extend emb/head rows.
    New key/value ids are initialised from a donor id (j mod 1024) plus small
    noise, so fresh ids start inside the working recall circuit instead of as
    random noise -- extending a working vocabulary rather than re-bootstrapping.
    From-scratch training at V>=16K historically never escaped the chance plateau
    (logs/vocab8192.log, logs/vocabcur.log).
    """
    ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    old = ck["model"]
    new = {k: (v.clone() if k not in ("emb.weight", "head.weight") else None)
           for k, v in new_model.state_dict().items()}
    # every non-emb/head parameter has identical shape -> straight copy
    for k in new:
        if new[k] is None:
            continue
        if k not in old or old[k].shape != new[k].shape:
            raise KeyError(f"cannot warm-start: parameter {k} missing/shape-mismatched")
        new[k] = old[k].clone()
    g = torch.Generator().manual_seed(seed)

    def build_rows(rows_old, std, off_new, span_new, off_old, span_old):
        out = torch.empty(span_new, rows_old.shape[1])
        donors = torch.arange(span_new) % span_old
        out[:] = rows_old[off_old + donors]
        out += torch.randn(out.shape, generator=g) * (noise * std)
        return out

    for pname in ("emb.weight", "head.weight"):
        rows_old, rows_new = old[pname], new_model.state_dict()[pname]
        std = rows_old.std().item()
        # keys: first old_NK ids keep their trained rows, the rest inherit donors
        rows_new[:old_NK] = rows_old[:old_NK]
        rows_new[old_NK:NK] = build_rows(rows_old, std, old_NK, NK - old_NK, 0, old_NK)
        # values (offset NK in the new layout, old_NK in the old one)
        rows_new[NK:NK + old_NV] = rows_old[old_NK:old_NK + old_NV]
        rows_new[NK + old_NV:NK + NV] = build_rows(rows_old, std, 0, NV - old_NV, old_NK, old_NV)
        # unused id + SEP
        rows_new[NK + NV] = rows_old[old_NK + old_NV]
        rows_new[NK + NV + 1] = rows_old[old_NK + old_NV + 1]
        new[pname] = rows_new
    new_model.load_state_dict(new)
    return new_model


# ---------------------------------------------------------------------------
# real-data protocol (same mining as experiments/eval_realdata.py)
# ---------------------------------------------------------------------------


class RealData:
    def __init__(self, corpus_path: str, VS: int = 4094):
        from collections import Counter
        self.VS = VS
        words = open(corpus_path, encoding="utf-8", errors="replace").read().lower().split()
        cnt = Counter(words)
        vocab = [w for w, _ in cnt.most_common(VS - 1)]
        wid = {w: i for i, w in enumerate(vocab)}
        big = Counter()
        for a, b in zip(words, words[1:]):
            if a in wid and b in wid:
                big[(wid[a], wid[b])] += 1
        succ = {}
        for (a, b), c in big.items():
            if c >= 4 and (a not in succ or succ[a][1] < c):
                succ[a] = (b, c)
        kv = {}
        ranked = sorted(succ.keys(), key=lambda k: -cnt[vocab[k]])
        lo, hi = int(len(ranked) * 0.02), int(len(ranked) * 0.95)
        for k in ranked[lo:hi]:
            v = succ[k][0]
            if v != k and cnt[vocab[v]] >= 10:
                kv[k] = v
        self.KVLIST = torch.tensor(sorted(kv.keys()))
        self.KVARR = torch.tensor([kv[int(k)] for k in self.KVLIST])
        # Zipfian position weights (cnt^0.75, as in realdata256k.py sample_pairs):
        # concentrating on frequent keys is what let the old run transition in
        # <400 steps; uniform sampling left rare-key embeddings gradient-starved
        # and the model stuck at 0.35 (measured 2026-11-02, realdata_v2 first run).
        w = torch.tensor([float(cnt[vocab[int(k)]]) ** 0.75 for k in self.KVLIST])
        self.KEY_W = w / w.sum()
        self.FILLER = torch.tensor([wid[w] for w in words if w in wid], dtype=torch.long)
        self.KVOCAB = VS
        self.SEP = 2 * VS
        self.V = 2 * VS + 2
        self.n_pairs = len(self.KVLIST)

    def sample_positions(self, B: int, P: int, g: torch.Generator, zipf: bool) -> torch.Tensor:
        """(B,P) distinct KVLIST positions; uniform or Zipf-weighted per row."""
        if not zipf:
            return torch.stack([torch.randperm(self.n_pairs, generator=g)[:P] for _ in range(B)])
        rows = []
        for _ in range(B):
            cand = torch.multinomial(self.KEY_W, 2 * P, replacement=True, generator=g).tolist()
            seen, out, j = set(), [], 0
            while len(out) < P and j < len(cand):
                k = cand[j]; j += 1
                if k not in seen:
                    seen.add(k); out.append(k)
            if len(out) < P:                      # exhausted weighted draws: fill uniformly
                rest = [i for i in torch.randperm(self.n_pairs, generator=g).tolist()
                        if i not in seen][: P - len(out)]
                out += rest
            rows.append(out)
        return torch.tensor(rows)

    def batch(self, P: int, filler: int, Q: int, seed: int, B: int = 1,
              overwrite_frac: float = 0.0, qpos_idx: torch.Tensor | None = None,
              zipf: bool = False):
        """Real-data store/query batch. qpos_idx: (Q,) indices into the P stored
        pairs (allows depth-stratified query sets); overwrite re-presents chosen
        keys with a DIFFERENT value token (another pair's value)."""
        g = torch.Generator().manual_seed(seed)
        P = min(P, self.n_pairs)
        perms = self.sample_positions(B, P, g, zipf)
        keys = self.KVLIST[perms]                       # (B,P) key token ids
        vals = self.KVARR[perms]                        # (B,P) value token ids
        n_upd = int(round(P * overwrite_frac))
        upd_idx = None
        upd_vals = None
        if n_upd:
            upd_idx = torch.stack([torch.randperm(P, generator=g)[:n_upd] for _ in range(B)])
            # update value = another pair's value; redraw the rare self-collisions
            shift = torch.stack([torch.randperm(P, generator=g)[:n_upd] for _ in range(B)])
            upd_vals = torch.gather(vals, 1, shift)
            same = upd_vals == torch.gather(vals, 1, upd_idx)
            while bool(same.any()):
                # rare: key's successor equals donor's successor -> bump by one token id
                flat = same.reshape(-1)
                upd_vals.reshape(-1)[flat] = (upd_vals.reshape(-1)[flat] + 1) % self.VS
                same = upd_vals == torch.gather(vals, 1, upd_idx)
        if qpos_idx is None:
            qpos_idx = torch.stack([torch.randperm(P, generator=g)[:Q] for _ in range(B)])
        elif qpos_idx.dim() == 1:
            qpos_idx = qpos_idx.unsqueeze(0).expand(B, -1).contiguous()
        Q = qpos_idx.shape[1]          # the actual query count (clamped to P)
        qkeys = torch.gather(keys, 1, qpos_idx)

        f = filler
        block = f + 2
        L = (P + n_upd) * block + 2 * Q
        tokens = torch.zeros(B, L, dtype=torch.long)
        fptr = torch.randint(0, max(1, len(self.FILLER) - f * (P + n_upd) - 1), (1,),
                             generator=g).item()
        for i in range(P):
            s = i * block
            if f:
                tokens[:, s:s + f] = self.FILLER[fptr:fptr + f]
                fptr += f
            tokens[:, s + f] = keys[:, i]
            tokens[:, s + f + 1] = vals[:, i] + self.KVOCAB
        if n_upd:
            upd_keys = torch.gather(keys, 1, upd_idx)
            for j in range(n_upd):
                s = (P + j) * block
                if f:
                    tokens[:, s:s + f] = self.FILLER[fptr:fptr + f]
                    fptr += f
                tokens[:, s + f] = upd_keys[:, j]
                tokens[:, s + f + 1] = upd_vals[:, j] + self.KVOCAB
        S2 = (P + n_upd) * block
        latest = vals.clone()
        if n_upd:
            latest.scatter_(1, upd_idx, upd_vals)
        gold = torch.gather(latest, 1, qpos_idx) + self.KVOCAB
        for q in range(Q):
            tokens[:, S2 + 2 * q] = qkeys[:, q]
            tokens[:, S2 + 2 * q + 1] = self.SEP
        key_pos = (torch.arange(P) * block + f).unsqueeze(0).expand(B, -1).contiguous()
        upd_key_pos = ((P + torch.arange(n_upd)) * block + f).unsqueeze(0).expand(B, -1).contiguous() \
            if n_upd else None
        # ANSWER AT THE QUERY-KEY POSITION (even offsets), matching the original
        # realdata256k.py training and eval_realdata.py: `labels[L-2*(n-j)]` is the
        # query-key token, not the SEP after it.  (The SYNTHETIC task answers at
        # the SEP -- mqar.py `ans_pos = arange(1, 2*n_queries, 2)` -- two different
        # conventions, both from the previous session; the warm checkpoint knows
        # the query-key one.  Measured: same tokens score 0.875 under the key
        # convention and 0.000 under the SEP convention.)
        ans_pos = (S2 + 2 * torch.arange(Q)).unsqueeze(0).expand(B, -1).contiguous()
        was_upd = torch.zeros(B, Q, dtype=torch.bool)
        if n_upd:
            m = torch.zeros(B, P, dtype=torch.bool)
            m.scatter_(1, upd_idx, True)
            was_upd = torch.gather(m, 1, qpos_idx)
        return {"tokens": tokens, "ans_pos": ans_pos, "gold": gold, "was_updated": was_upd,
                "key_pos": key_pos, "upd_key_pos": upd_key_pos, "n_pairs": P, "n_upd": n_upd,
                "L": L}


def stratified_qpos(P: int, Q: int, strata: int = 4, seed: int = 0):
    """Query indices stratified over store depth: each stratum gets Q/strata
    queries drawn from pairs stored at that depth fraction of the sequence."""
    g = torch.Generator().manual_seed(seed)
    per = Q // strata
    out = []
    bounds = torch.linspace(0, P, strata + 1).long()
    for s in range(strata):
        lo, hi = int(bounds[s]), int(bounds[s + 1])
        out.append(lo + torch.randperm(hi - lo, generator=g)[:per])
    return torch.cat(out)

	"""Shared evaluation/training protocol for the capacity, forgetting and real-data goals.

	Single source of truth: the acceptance tests (tests/test_capacity.py,
	tests/test_forgetting.py, tests/test_realtext.py) and the training scripts
	(experiments/train_capacity.py, experiments/train_forgetting.py,
	experiments/realdata_v2.py) all build sequences and measure accuracy through
	this module, so "train until the test passes" cannot drift from what is tested.

	Synthetic task (capacity + forgetting):
	vocab keys [0, NK) values [NK, NK+NV) SEP = NK+NV+1 V = NK+NV+2
	sequence per row (filler f tokens before each pair; f=0 -> dense MQAR):
	[f k1 v1] [f k2 v2] ... [f kP vP] store pass
	[f ku1 vu1'] ... [f kuM vuM'] update pass (overwrites)
	[q1 SEP] [q2 SEP] ... [qQ SEP] query pass
	A query on an updated key must return the UPDATED value (last write wins).
	The answer is read at the SEP token of each query (causal).

	Real-data task: wikitext-103 mined key->bigram-successor pairs embedded in a
	real corpus filler stream (same mining as experiments/eval_realdata.py).
	"""
	from __future__ import annotations

	import os
	import time
	from dataclasses import dataclass

	import torch

	HERE = os.path.dirname(os.path.abspath(__file__))
	ROOT = os.path.join(HERE, "..")

	# ---------------------------------------------------------------------------
	# synthetic generator
	# ---------------------------------------------------------------------------


	@dataclass
	class SynthCfg:
	NK: int = 16384 # distinct key tokens available
	NV: int = 16384 # distinct value tokens available
	d_model: int = 64
	n_heads: int = 2
	d_k: int = 32
	d_v: int = 32
	addr_scale: float = 64.0 # (2*64+3)^2 = 17,161 cells for 8,192 keys
	retention_init: float = 9.0

	@property
	def V(self) -> int:
	return self.NK + self.NV + 2

	@property
	def SEP(self) -> int:
	return self.NK + self.NV + 1


	def synth_batch(
	B: int,
	n_pairs: int,
	n_queries: int,
	cfg: SynthCfg,
	filler: int = 0,
	overwrite_frac: float = 0.0,
	generator: torch.Generator \| None = None,
	qpos: torch.Tensor \| None = None,
	updated_idx: torch.Tensor \| None = None,
	kw: dict \| None = None,
	):
	"""Build a capacity/forgetting batch. See module docstring for the layout.

	Returns dict with:
	tokens (B, L) long the sequence
	ans_pos (B, Q) long SEP positions; answer read there
	gold (B, Q) long value token id that must be produced
	was_updated (B, Q) bool query key was overwritten -> updated gold
	key_pos (B, P) long first-pass key token positions (spread loss)
	upd_key_pos (B, M) long \| None update-pass key token positions
	n_pairs, n_upd, L
	"""
	g = generator
	NK, NV = cfg.NK, cfg.NV
	if n_pairs > NK:
	raise ValueError("n_pairs > key vocab")
	n_upd = int(round(n_pairs * overwrite_frac))
	if n_pairs + n_upd > NV:
	raise ValueError("not enough distinct values")
	dev = "cpu" # build on CPU; caller moves to device

	# distinct keys per row; `include_keys` are guaranteed members of the stored
	# set (placed FIRST), and the rest are sampled from the complement, so the
	# batch is consistent by construction -- no duplicate stores of the forced
	# keys with conflicting values (the token-surgery approach produced exactly
	# that and wrecked the model: census 9 -> 713 errors, polish_weakkeys v1).
	include_keys = None
	if "include_keys" in (kw or {}):
	include_keys = kw["include_keys"]
	if include_keys.dim() == 1:
	include_keys = include_keys.unsqueeze(0).expand(B, -1)
	if include_keys is not None:
	nk_inc = include_keys.shape[1]
	assert n_pairs >= nk_inc
	rows = []
	for r in range(B):
	rest = torch.randperm(NK, generator=g)
	rest = rest[torch.isin(rest, include_keys[r], invert=True)][: n_pairs - nk_inc]
	rows.append(torch.cat([include_keys[r], rest]))
	keys = torch.stack(rows)
	else:
	keys = torch.stack([torch.randperm(NK, generator=g)[:n_pairs] for _ in range(B)])
	# distinct values: n_pairs originals + n_upd updates, all different in a row
	vals = torch.stack([torch.randperm(NV, generator=g)[: n_pairs + n_upd] for _ in range(B)])
	orig_vals, upd_vals = vals[:, :n_pairs], vals[:, n_pairs:]

	if updated_idx is None:
	# which stored keys get overwritten (row-specific)
	updated_idx = torch.stack([torch.randperm(n_pairs, generator=g)[:n_upd] for _ in range(B)])
	if qpos is None:
	# DISTINCT query keys (sample without replacement). Sampling with
	# replacement creates duplicate queries at large Q (measured: 9 of 512
	# at Q=512/NK=16,384) -- not part of the published protocol semantics
	# ("every query has a unique answer", mqar.py); distinct sampling is
	# the clean generalization of the published small-Q generator.
	qpos = torch.stack([torch.randperm(n_pairs, generator=g)[:n_queries] for _ in range(B)])
	else:
	if qpos.dim() == 1:
	qpos = qpos.unsqueeze(0).expand(B, -1).contiguous()
	n_queries = qpos.shape[1] # explicit query set defines Q
	qkeys = torch.gather(keys, 1, qpos)
	upd_mask = torch.zeros(B, n_pairs, dtype=torch.bool)
	upd_mask.scatter_(1, updated_idx, True)
	was_updated = torch.gather(upd_mask, 1, qpos)
	# latest value for every stored key: original, replaced where updated
	latest = orig_vals.clone()
	latest.scatter_(1, updated_idx, upd_vals)
	gold = torch.gather(latest, 1, qpos) + NK

	f = filler
	block = f + 2
	L = (n_pairs + n_upd) * block + 2 * n_queries

	tokens = torch.zeros(B, L, dtype=torch.long)
	# filler ids: uniform from the complement of the row's stored keys
	if f > 0:
	need = f * (n_pairs + n_upd)
	filler_ids = torch.zeros(B, need, dtype=torch.long)
	keyset = torch.zeros(B, NK, dtype=torch.bool)
	keyset.scatter_(1, keys, True)
	done = torch.zeros(B, need, dtype=torch.bool)
	while not bool(done.all()):
	cand = torch.randint(0, NK, (B, need), generator=g)
	ok = ~torch.gather(keyset, 1, cand) & ~done
	filler_ids = torch.where(ok, cand, filler_ids)
	done = done \| ok
	# store pass
	for i in range(n_pairs):
	s = i * block
	if f:
	tokens[:, s:s + f] = filler_ids[:, i * f:(i + 1) * f]
	tokens[:, s + f] = keys[:, i]
	tokens[:, s + f + 1] = orig_vals[:, i] + NK
	# update pass: same key, new value, later in the sequence
	upd_tokens = torch.gather(keys, 1, updated_idx)
	for j in range(n_upd):
	s = (n_pairs + j) * block
	if f:
	tokens[:, s:s + f] = filler_ids[:, (n_pairs + j) * f:(n_pairs + j + 1) * f]
	tokens[:, s + f] = upd_tokens[:, j]
	tokens[:, s + f + 1] = upd_vals[:, j] + NK
	# query pass
	S2 = (n_pairs + n_upd) * block
	for q in range(n_queries):
	tokens[:, S2 + 2 * q] = qkeys[:, q]
	tokens[:, S2 + 2 * q + 1] = cfg.SEP

	key_pos = (torch.arange(n_pairs) * block + f).unsqueeze(0).expand(B, -1).contiguous()
	upd_key_pos = ((n_pairs + torch.arange(n_upd)) * block + f).unsqueeze(0).expand(B, -1).contiguous() \
	if n_upd else None
	ans_pos = (S2 + 2 * torch.arange(n_queries) + 1).unsqueeze(0).expand(B, -1).contiguous()
	return {
	"tokens": tokens, "ans_pos": ans_pos, "gold": gold, "was_updated": was_updated,
	"key_pos": key_pos, "upd_key_pos": upd_key_pos, "n_pairs": n_pairs,
	"n_upd": n_upd, "L": L,
	}


	# ---------------------------------------------------------------------------
	# model helpers
	# ---------------------------------------------------------------------------


	def build_model(cfg: SynthCfg, n_layers: int = 1):
	from mulattice.models import MQARModel, MixerCfg
	mcfg = MixerCfg(d_model=cfg.d_model, n_heads=cfg.n_heads, d_k=cfg.d_k, d_v=cfg.d_v,
	addr_scale=cfg.addr_scale, learn_addr_scale=False, bound_address=True,
	read_tie=False, retention_init=cfg.retention_init)
	return MQARModel(cfg.V, kind="mulattice", n_layers=n_layers, cfg=mcfg)


	def logits_at(model, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
	"""Forward the trunk, then apply the head ONLY at `positions` (B,Q).

	Avoids materialising (B,T,V) logits; at T=131K with V=32,770 that is 17GB.
	"""
	x = model.emb(tokens)
	for blk in model.blocks:
	x = blk(x, read_mode="all")
	x = model.norm(x)
	B = x.shape[0]
	h = x[torch.arange(B, device=x.device).unsqueeze(1), positions]
	return model.head(h) # (B,Q,V)


	@torch.no_grad()
	def eval_batch(model, batch, device) -> dict:
	"""Accuracy of `model` on a synth_batch, split by updated/untouched queries."""
	model.eval()
	pred = logits_at(model, batch["tokens"].to(device), batch["ans_pos"].to(device)).argmax(-1).cpu()
	gold = batch["gold"]
	ok = pred == gold
	res = {"acc": ok.float().mean().item(), "n": ok.numel(),
	"correct": int(ok.sum())}
	up = batch["was_updated"]
	if bool(up.any()):
	res["acc_updated"] = ok[up].float().mean().item()
	if bool((~up).any()):
	res["acc_untouched"] = ok[~up].float().mean().item()
	model.train()
	return res


	def scan_info(model) -> str:
	parts = []
	for i, blk in enumerate(model.blocks):
	sc = blk.mixer.scan
	parts.append(f"L{i} cells={sc.last_n_cells} rounds={getattr(sc, 'last_rounds', '?')}"
	f" path={getattr(sc, 'last_path', '?')}")
	return " ".join(parts)


	def force_fast_scan(model):
	for blk in model.blocks:
	blk.mixer.scan.crossover = 1e9


	def key_addresses(model, tokens, key_pos):
	"""Write addresses (B,N,2) of the tokens at `key_pos` (B,N).

	Uses the mixer's LIVE addr_scale (it can be ramped mid-run), not the
	construction-time cfg value."""
	mx = model.blocks[0].mixer
	xc = mx.conv(model.blocks[0].n1(model.emb(tokens)))
	kaddr = mx.cfg.addr_scale * torch.tanh(mx.Wa(xc))
	ar = torch.arange(tokens.shape[0], device=tokens.device).unsqueeze(1)
	return kaddr[ar, key_pos]


	def spread_loss(model, tokens, key_pos, cfg: SynthCfg, gate_subsample: int = 2048):
	"""Kernel-D gated spread loss on the TRUE key-token write addresses.

	Same definition as mulattice.addressing.gated_spread_loss, but the
	pairwise write-overlap gate is computed on a subsample when N is large
	(it is (B,N,N); at N=8,192 that is 67M entries per row)."""
	addr = key_addresses(model, tokens, key_pos)
	B, N, _ = addr.shape
	addr_gate = addr[:, torch.randperm(N, device=addr.device)[:gate_subsample]] if N > gate_subsample else addr
	return _gated_spread(addr, addr_gate, float(model.blocks[0].mixer.cfg.addr_scale))


	def _gated_spread(addr, addr_gate, scale):
	from mulattice.addressing import write_overlap
	from mulattice.kernel import phi_2d
	B, N, _ = addr.shape
	gate = (write_overlap(addr_gate) - 0.05).clamp(0.0, 1.0).detach()
	lam = 0.3 + 0.7 * gate
	w, z = phi_2d(addr)
	size = int(2 * scale + 3)
	off = int(scale) + 1
	gi = (z + off).clamp(0, size - 1)
	idx = (gi[..., 0] * size + gi[..., 1]).reshape(B, -1).to(torch.int64)
	mass = (w / w.sum(dim=(-2, -1), keepdim=True)).reshape(B, -1)
	hist = torch.zeros(B, size * size, device=addr.device, dtype=addr.dtype)
	for b in range(B):
	hist[b].index_add_(0, idx[b], mass[b])
	p = hist / hist.sum(dim=1, keepdim=True).clamp_min(1e-9)
	entropy = -(p * (p + 1e-12).log()).sum(dim=1)
	return (lam * (-entropy)).mean()


	# ---------------------------------------------------------------------------
	# vocab surgery: extend the proven 1024-key model to the full key vocabulary
	# ---------------------------------------------------------------------------


	def extend_from_checkpoint(new_model, ckpt_path: str, NK: int = 16384, NV: int = 16384,
	old_NK: int = 1024, old_NV: int = 1024,
	noise: float = 0.05, seed: int = 0):
	"""Copy the proven model's mixer into `new_model`; extend emb/head rows.

	New key/value ids are initialised from a donor id (j mod 1024) plus small
	noise, so fresh ids start inside the working recall circuit instead of as
	random noise -- extending a working vocabulary rather than re-bootstrapping.
	From-scratch training at V>=16K historically never escaped the chance plateau
	(logs/vocab8192.log, logs/vocabcur.log).
	"""
	ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
	old = ck["model"]
	new = {k: (v.clone() if k not in ("emb.weight", "head.weight") else None)
	for k, v in new_model.state_dict().items()}
	# every non-emb/head parameter has identical shape -> straight copy
	for k in new:
	if new[k] is None:
	continue
	if k not in old or old[k].shape != new[k].shape:
	raise KeyError(f"cannot warm-start: parameter {k} missing/shape-mismatched")
	new[k] = old[k].clone()
	g = torch.Generator().manual_seed(seed)

	def build_rows(rows_old, std, off_new, span_new, off_old, span_old):
	out = torch.empty(span_new, rows_old.shape[1])
	donors = torch.arange(span_new) % span_old
	out[:] = rows_old[off_old + donors]
	out += torch.randn(out.shape, generator=g) * (noise * std)
	return out

	for pname in ("emb.weight", "head.weight"):
	rows_old, rows_new = old[pname], new_model.state_dict()[pname]
	std = rows_old.std().item()
	# keys: first old_NK ids keep their trained rows, the rest inherit donors
	rows_new[:old_NK] = rows_old[:old_NK]
	rows_new[old_NK:NK] = build_rows(rows_old, std, old_NK, NK - old_NK, 0, old_NK)
	# values (offset NK in the new layout, old_NK in the old one)
	rows_new[NK:NK + old_NV] = rows_old[old_NK:old_NK + old_NV]
	rows_new[NK + old_NV:NK + NV] = build_rows(rows_old, std, 0, NV - old_NV, old_NK, old_NV)
	# unused id + SEP
	rows_new[NK + NV] = rows_old[old_NK + old_NV]
	rows_new[NK + NV + 1] = rows_old[old_NK + old_NV + 1]
	new[pname] = rows_new
	new_model.load_state_dict(new)
	return new_model


	# ---------------------------------------------------------------------------
	# real-data protocol (same mining as experiments/eval_realdata.py)
	# ---------------------------------------------------------------------------


	class RealData:
	def __init__(self, corpus_path: str, VS: int = 4094):
	from collections import Counter
	self.VS = VS
	words = open(corpus_path, encoding="utf-8", errors="replace").read().lower().split()
	cnt = Counter(words)
	vocab = [w for w, _ in cnt.most_common(VS - 1)]
	wid = {w: i for i, w in enumerate(vocab)}
	big = Counter()
	for a, b in zip(words, words[1:]):
	if a in wid and b in wid:
	big[(wid[a], wid[b])] += 1
	succ = {}
	for (a, b), c in big.items():
	if c >= 4 and (a not in succ or succ[a][1] < c):
	succ[a] = (b, c)
	kv = {}
	ranked = sorted(succ.keys(), key=lambda k: -cnt[vocab[k]])
	lo, hi = int(len(ranked) * 0.02), int(len(ranked) * 0.95)
	for k in ranked[lo:hi]:
	v = succ[k][0]
	if v != k and cnt[vocab[v]] >= 10:
	kv[k] = v
	self.KVLIST = torch.tensor(sorted(kv.keys()))
	self.KVARR = torch.tensor([kv[int(k)] for k in self.KVLIST])
	# Zipfian position weights (cnt^0.75, as in realdata256k.py sample_pairs):
	# concentrating on frequent keys is what let the old run transition in
	# <400 steps; uniform sampling left rare-key embeddings gradient-starved
	# and the model stuck at 0.35 (measured 2026-11-02, realdata_v2 first run).
	w = torch.tensor([float(cnt[vocab[int(k)]]) ** 0.75 for k in self.KVLIST])
	self.KEY_W = w / w.sum()
	self.FILLER = torch.tensor([wid[w] for w in words if w in wid], dtype=torch.long)
	self.KVOCAB = VS
	self.SEP = 2 * VS
	self.V = 2 * VS + 2
	self.n_pairs = len(self.KVLIST)

	def sample_positions(self, B: int, P: int, g: torch.Generator, zipf: bool) -> torch.Tensor:
	"""(B,P) distinct KVLIST positions; uniform or Zipf-weighted per row."""
	if not zipf:
	return torch.stack([torch.randperm(self.n_pairs, generator=g)[:P] for _ in range(B)])
	rows = []
	for _ in range(B):
	cand = torch.multinomial(self.KEY_W, 2 * P, replacement=True, generator=g).tolist()
	seen, out, j = set(), [], 0
	while len(out) < P and j < len(cand):
	k = cand[j]; j += 1
	if k not in seen:
	seen.add(k); out.append(k)
	if len(out) < P: # exhausted weighted draws: fill uniformly
	rest = [i for i in torch.randperm(self.n_pairs, generator=g).tolist()
	if i not in seen][: P - len(out)]
	out += rest
	rows.append(out)
	return torch.tensor(rows)

	def batch(self, P: int, filler: int, Q: int, seed: int, B: int = 1,
	overwrite_frac: float = 0.0, qpos_idx: torch.Tensor \| None = None,
	zipf: bool = False):
	"""Real-data store/query batch. qpos_idx: (Q,) indices into the P stored
	pairs (allows depth-stratified query sets); overwrite re-presents chosen
	keys with a DIFFERENT value token (another pair's value)."""
	g = torch.Generator().manual_seed(seed)
	P = min(P, self.n_pairs)
	perms = self.sample_positions(B, P, g, zipf)
	keys = self.KVLIST[perms] # (B,P) key token ids
	vals = self.KVARR[perms] # (B,P) value token ids
	n_upd = int(round(P * overwrite_frac))
	upd_idx = None
	upd_vals = None
	if n_upd:
	upd_idx = torch.stack([torch.randperm(P, generator=g)[:n_upd] for _ in range(B)])
	# update value = another pair's value; redraw the rare self-collisions
	shift = torch.stack([torch.randperm(P, generator=g)[:n_upd] for _ in range(B)])
	upd_vals = torch.gather(vals, 1, shift)
	same = upd_vals == torch.gather(vals, 1, upd_idx)
	while bool(same.any()):
	# rare: key's successor equals donor's successor -> bump by one token id
	flat = same.reshape(-1)
	upd_vals.reshape(-1)[flat] = (upd_vals.reshape(-1)[flat] + 1) % self.VS
	same = upd_vals == torch.gather(vals, 1, upd_idx)
	if qpos_idx is None:
	qpos_idx = torch.stack([torch.randperm(P, generator=g)[:Q] for _ in range(B)])
	elif qpos_idx.dim() == 1:
	qpos_idx = qpos_idx.unsqueeze(0).expand(B, -1).contiguous()
	Q = qpos_idx.shape[1] # the actual query count (clamped to P)
	qkeys = torch.gather(keys, 1, qpos_idx)

	f = filler
	block = f + 2
	L = (P + n_upd) * block + 2 * Q
	tokens = torch.zeros(B, L, dtype=torch.long)
	fptr = torch.randint(0, max(1, len(self.FILLER) - f * (P + n_upd) - 1), (1,),
	generator=g).item()
	for i in range(P):
	s = i * block
	if f:
	tokens[:, s:s + f] = self.FILLER[fptr:fptr + f]
	fptr += f
	tokens[:, s + f] = keys[:, i]
	tokens[:, s + f + 1] = vals[:, i] + self.KVOCAB
	if n_upd:
	upd_keys = torch.gather(keys, 1, upd_idx)
	for j in range(n_upd):
	s = (P + j) * block
	if f:
	tokens[:, s:s + f] = self.FILLER[fptr:fptr + f]
	fptr += f
	tokens[:, s + f] = upd_keys[:, j]
	tokens[:, s + f + 1] = upd_vals[:, j] + self.KVOCAB
	S2 = (P + n_upd) * block
	latest = vals.clone()
	if n_upd:
	latest.scatter_(1, upd_idx, upd_vals)
	gold = torch.gather(latest, 1, qpos_idx) + self.KVOCAB
	for q in range(Q):
	tokens[:, S2 + 2 * q] = qkeys[:, q]
	tokens[:, S2 + 2 * q + 1] = self.SEP
	key_pos = (torch.arange(P) * block + f).unsqueeze(0).expand(B, -1).contiguous()
	upd_key_pos = ((P + torch.arange(n_upd)) * block + f).unsqueeze(0).expand(B, -1).contiguous() \
	if n_upd else None
	# ANSWER AT THE QUERY-KEY POSITION (even offsets), matching the original
	# realdata256k.py training and eval_realdata.py: `labels[L-2*(n-j)]` is the
	# query-key token, not the SEP after it. (The SYNTHETIC task answers at
	# the SEP -- mqar.py `ans_pos = arange(1, 2*n_queries, 2)` -- two different
	# conventions, both from the previous session; the warm checkpoint knows
	# the query-key one. Measured: same tokens score 0.875 under the key
	# convention and 0.000 under the SEP convention.)
	ans_pos = (S2 + 2 * torch.arange(Q)).unsqueeze(0).expand(B, -1).contiguous()
	was_upd = torch.zeros(B, Q, dtype=torch.bool)
	if n_upd:
	m = torch.zeros(B, P, dtype=torch.bool)
	m.scatter_(1, upd_idx, True)
	was_upd = torch.gather(m, 1, qpos_idx)
	return {"tokens": tokens, "ans_pos": ans_pos, "gold": gold, "was_updated": was_upd,
	"key_pos": key_pos, "upd_key_pos": upd_key_pos, "n_pairs": P, "n_upd": n_upd,
	"L": L}


	def stratified_qpos(P: int, Q: int, strata: int = 4, seed: int = 0):
	"""Query indices stratified over store depth: each stratum gets Q/strata
	queries drawn from pairs stored at that depth fraction of the sequence."""
	g = torch.Generator().manual_seed(seed)
	per = Q // strata
	out = []
	bounds = torch.linspace(0, P, strata + 1).long()
	for s in range(strata):
	lo, hi = int(bounds[s]), int(bounds[s + 1])
	out.append(lo + torch.randperm(hi - lo, generator=g)[:per])
	return torch.cat(out)