spb/wpai0001_uqo

1"""Reusable pieces of the Tier-0 transplant experiments: split construction, activation-row collection,2condition evaluation with grafts installed, supervised adapter baseline and distillation baseline."""3from __future__ import annotations45import copy6import time7from dataclasses import dataclass89import numpy as np10import torch11import torch.nn.functional as F1213from wp_ai_0001.alignment.maps import AffineMap, fit_ridge14from wp_ai_0001.config import OpSeqTaskConfig15from wp_ai_0001.datasets import opseq16from wp_ai_0001.datasets.common import Split, Vocab, encode_split17from wp_ai_0001.evaluation.metrics import EvalResult, evaluate18from wp_ai_0001.models.hooks import Capture19from wp_ai_0001.models.toy_transformer import ToyTransformer2021OMEGA = "qua"22MAX_SEQ = 24232425# ----------------------------------------------------------------------------------------------26# Data27# ----------------------------------------------------------------------------------------------28def omega_read_mask(split: Split, vocab: Vocab, max_len: int = MAX_SEQ) -> np.ndarray:29    """[N, T] bool: True at read positions whose next emitted token is the result of an Ω step30    (trace_stride == 1 only; for stride > 1, the step containing Ω)."""31    enc = encode_split(split, vocab, max_len)32    mask = np.zeros_like(enc["ids"], dtype=bool)33    for i, e in enumerate(split.examples):34        ops = list(e.prompt[2:-1])35        stride = e.meta.get("stride", 1)36        pl = enc["prompt_len"][i]37        # visible step j covers ops[(j*stride) : (j+1)*stride] (last step may be shorter)38        n_vis = e.n_answer39        for j in range(n_vis):40            lo = j * stride41            hi = min(len(ops), (j + 1) * stride) if j < n_vis - 1 else len(ops)42            if OMEGA in ops[lo:hi]:43                mask[i, pl - 1 + j] = True44    return mask454647def omega_input_values(split: Split) -> list[list[int]]:48    """For each item, the list of input values to Ω occurrences (value before applying Ω)."""49    out = []50    for e in split.examples:51        ops = list(e.prompt[2:-1])52        v = int(e.tokens[1])53        vals = []54        p = int(e.meta.get("modulus", 11)) if "modulus" in e.meta else 1155        for op in ops:56            if op == OMEGA:57                vals.append(v)58            v = opseq.OPERATOR_TABLE[op](v, p)59        out.append(vals)60    return out616263def build_eval_splits(task: OpSeqTaskConfig, exclude: set[str], base_seed: int, n_eval: int, n_align: int) -> dict[str, Split]:64    s = base_seed65    sp: dict[str, Split] = {}66    used = set(exclude)67    sp["align_target"] = opseq.generate(task, n_align, s + 1, "align_target", mode="train", require_operator=OMEGA, exclude_hashes=used)68    used |= sp["align_target"].item_hashes()69    sp["align_unrelated"] = opseq.generate(task, n_align, s + 2, "align_unrelated", mode="train", forbid_operator=OMEGA, exclude_hashes=used)70    used |= sp["align_unrelated"].item_hashes()71    sp["target_iid"] = opseq.generate(task, n_eval, s + 10, "target_iid", mode="train", require_operator=OMEGA, exclude_hashes=used)72    used |= sp["target_iid"].item_hashes()73    sp["control_iid"] = opseq.generate(task, n_eval, s + 11, "control_iid", mode="train", forbid_operator=OMEGA, exclude_hashes=used)74    used |= sp["control_iid"].item_hashes()75    sp["comp_heldout"] = opseq.generate(task, n_eval, s + 12, "comp_heldout", mode="comp", min_len=2, exclude_hashes=used)76    used |= sp["comp_heldout"].item_hashes()77    sp["len_ood"] = opseq.generate(task, max(200, n_eval // 2), s + 13, "len_ood", mode="any", min_len=6, max_len=7, exclude_hashes=used)78    for v in sp.values():79        opseq.validate(v, task)80    return sp818283# ----------------------------------------------------------------------------------------------84# Activation rows85# ----------------------------------------------------------------------------------------------86@dataclass87class Rows:88    X: np.ndarray  # [M, d]89    item: np.ndarray  # [M] item index90    step: np.ndarray  # [M] visible step index91    is_omega: np.ndarray  # [M] bool92    omega_input: np.ndarray  # [M] input value to Ω for Ω steps, -1 otherwise939495@torch.no_grad()96def collect_rows(model: ToyTransformer, split: Split, vocab: Vocab, hook: str, device: torch.device, bs: int = 512) -> Rows:97    """Teacher-forced activations at every read position of every item."""98    enc = encode_split(split, vocab, MAX_SEQ)99    omask = omega_read_mask(split, vocab)100    oin = omega_input_values(split)101    X, item, step, isom, oval = [], [], [], [], []102    for i0 in range(0, len(split), bs):103        ids = torch.from_numpy(enc["ids"][i0 : i0 + bs]).to(device)104        cap = Capture()105        model.hooks.add(hook, cap)106        try:107            model(ids)108        finally:109            model.hooks.remove(hook, cap)110        act = cap.store[hook].float().cpu().numpy()111        for r in range(ids.shape[0]):112            i = i0 + r113            pl, na = int(enc["prompt_len"][i]), int(enc["n_answer"][i])114            k_om = 0115            for j in range(na):116                p = pl - 1 + j117                X.append(act[r, p])118                item.append(i)119                step.append(j)120                om = bool(omask[i, p])121                isom.append(om)122                if om:123                    oval.append(oin[i][k_om] if k_om < len(oin[i]) else -1)124                    k_om += 1125                else:126                    oval.append(-1)127    return Rows(np.stack(X), np.array(item), np.array(step), np.array(isom), np.array(oval))128129130# ----------------------------------------------------------------------------------------------131# Evaluation with grafts installed132# ----------------------------------------------------------------------------------------------133@torch.no_grad()134def evaluate_condition(recipient: ToyTransformer, hook: str | None, graft, splits: dict[str, Split], vocab: Vocab, device: torch.device, names: list[str]) -> dict[str, EvalResult]:135    out = {}136    if hook is not None and graft is not None:137        recipient.hooks.add(hook, graft)138    try:139        for n in names:140            if graft is not None and getattr(graft, "omega_positions", None) is not None or (graft is not None and graft.spec.targeting == "omega"):141                graft.omega_positions = torch.from_numpy(omega_read_mask(splits[n], vocab)).to(device)142            out[n] = evaluate(recipient, splits[n], vocab, device)143    finally:144        if hook is not None and graft is not None:145            recipient.hooks.remove(hook, graft)146    return out147148149def omega_step_acc(split: Split, res: EvalResult) -> np.ndarray:150    """Per item: accuracy of the first Ω-step token (NaN for items without Ω)."""151    out = np.full(len(split.examples), np.nan)152    for i, (e, pred) in enumerate(zip(split.examples, res.predictions)):153        ops = list(e.prompt[2:-1])154        stride = e.meta.get("stride", 1)155        if OMEGA in ops:156            j = ops.index(OMEGA) // stride157            out[i] = float(j < len(pred) and pred[j] == e.answers[j])158    return out159160161# ----------------------------------------------------------------------------------------------162# Supervised baselines (conventional learning, matched data)163# ----------------------------------------------------------------------------------------------164def train_supervised_adapter(recipient: ToyTransformer, hook: str, split: Split, vocab: Vocab, device: torch.device, steps: int = 300, lr: float = 1e-2, seed: int = 0,165                             omega_only: bool = False) -> AffineMap:166    """Affine adapter on the residual rows at the hook (same parameter count as the ridge adapter), trained by167    gradient descent on the recipient's cross-entropy for answer tokens of ``split`` (labels, no donor)."""168    torch.manual_seed(seed)169    d = recipient.cfg.d_model170    W = torch.nn.Parameter(torch.eye(d, device=device))171    b = torch.nn.Parameter(torch.zeros(d, device=device))172    opt = torch.optim.Adam([W, b], lr=lr)173    enc = encode_split(split, vocab, MAX_SEQ)174    omask = torch.from_numpy(omega_read_mask(split, vocab)).to(device)175    ids_all = torch.from_numpy(enc["ids"]).to(device)176    lm_all = torch.from_numpy(enc["loss_mask"]).to(device)177    if omega_only:178        lm_all = lm_all * omask.float()179180    def fn(x, ctx):181        return x @ W + b182183    recipient.hooks.add(hook, fn)184    g = torch.Generator().manual_seed(seed)185    try:186        for _ in range(steps):187            idx = torch.randint(0, len(ids_all), (256,), generator=g).to(device)188            ids, m = ids_all[idx], lm_all[idx]189            logits = recipient(ids[:, :-1])190            ce = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), ids[:, 1:].reshape(-1), reduction="none").view_as(m[:, :-1])191            loss = (ce * m[:, :-1]).sum() / m[:, :-1].sum().clamp_min(1)192            opt.zero_grad()193            loss.backward()194            opt.step()195    finally:196        recipient.hooks.remove(hook, fn)197    return AffineMap("supervised_adapter", W.detach().cpu().numpy().astype(np.float64), b.detach().cpu().numpy().astype(np.float64), {"steps": steps, "lr": lr, "n_fit": len(split)})198199200def distill(recipient: ToyTransformer, donor: ToyTransformer, split: Split, vocab: Vocab, device: torch.device, steps: int = 300, lr: float = 3e-4, seed: int = 0) -> ToyTransformer:201    """Conventional distillation baseline: fine-tune a COPY of the recipient on the donor's next-token202    distributions at answer positions of ``split`` (same data as the alignment set)."""203    student = copy.deepcopy(recipient).to(device)204    student.train()205    torch.manual_seed(seed)206    opt = torch.optim.AdamW(student.parameters(), lr=lr, weight_decay=0.01)207    enc = encode_split(split, vocab, MAX_SEQ)208    ids_all = torch.from_numpy(enc["ids"]).to(device)209    lm_all = torch.from_numpy(enc["loss_mask"]).to(device)210    g = torch.Generator().manual_seed(seed)211    for _ in range(steps):212        idx = torch.randint(0, len(ids_all), (256,), generator=g).to(device)213        ids, m = ids_all[idx], lm_all[idx][:, :-1]214        with torch.no_grad():215            t = F.log_softmax(donor(ids[:, :-1]), -1)216        s = F.log_softmax(student(ids[:, :-1]), -1)217        kl = (t.exp() * (t - s)).sum(-1)218        loss = (kl * m).sum() / m.sum().clamp_min(1)219        opt.zero_grad()220        loss.backward()221        opt.step()222    student.eval()223    return student224225226def fit_adapter_from_donor(rec_rows: np.ndarray, mapped_donor_rows: np.ndarray, lam: float = 1e-2) -> AffineMap:227    """Donor-free interface: predicts the mapped donor row from the recipient's own row (ridge)."""228    m = fit_ridge(rec_rows, mapped_donor_rows, lam=lam, family="adapter_ridge")229    return m230