job-code/run_experiment.py

#!/usr/bin/env python3
"""Independent reproduction of the estimator in Proposition 3.1.

The model is the output of the paper's constructive two-layer disentangled
transformer: layer one copies the k context tokens and layer two attends to
prior transitions with an exponentiated Hamming-overlap score.  Training the
free beta/kappa parameters is therefore an exact, interpretable realization
of the claimed circuit rather than a generic sequence-model proxy.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import random
import time
from pathlib import Path

import numpy as np
import torch
from torch import Tensor


def sample_chain(k: int, vocab: int, length: int, alpha: float, rng: np.random.Generator):
    nctx = vocab**k
    transition = rng.dirichlet(np.full(vocab, alpha), size=nctx)
    seq = [int(rng.integers(vocab)) for _ in range(k)]
    for _ in range(k, length):
        idx = 0
        for token in seq[-k:]:
            idx = idx * vocab + token
        seq.append(int(rng.choice(vocab, p=transition[idx])))
    return seq, transition


def context_index(context: list[int], vocab: int) -> int:
    out = 0
    for token in context:
        out = out * vocab + token
    return out


def predict(seq: list[int], t: int, beta: Tensor, vocab: int, log_bos_mass: Tensor | None) -> Tensor:
    """Predict x[t] from transitions ending before t (strictly causal)."""
    k = beta.numel()
    query = torch.tensor(seq[t-k:t], dtype=beta.dtype, device=beta.device)
    contexts = torch.tensor([seq[s-k:s] for s in range(k, t)], dtype=beta.dtype, device=beta.device)
    outcomes = torch.tensor(seq[k:t], dtype=torch.long, device=beta.device)
    matches = contexts.eq(query).to(beta.dtype)
    weights = torch.exp(matches @ beta)
    counts = torch.zeros(vocab, dtype=beta.dtype, device=beta.device).scatter_add(0, outcomes, weights)
    total = weights.sum()
    if log_bos_mass is not None:
        mass = torch.exp(log_bos_mass)
        counts = counts + mass / vocab
        total = total + mass
    return counts / total


def bayes_predict(seq: list[int], t: int, k: int, vocab: int, alpha: float, device: torch.device) -> Tensor:
    query = seq[t-k:t]
    counts = torch.full((vocab,), alpha, dtype=torch.float64, device=device)
    for s in range(k, t):
        if seq[s-k:s] == query:
            counts[seq[s]] += 1
    return counts / counts.sum()


def episode_loss(seq: list[int], beta: Tensor, vocab: int, log_bos_mass: Tensor | None, start: int) -> Tensor:
    losses = []
    for t in range(max(beta.numel() + 1, start), len(seq)):
        losses.append(-torch.log(predict(seq, t, beta, vocab, log_bos_mass)[seq[t]].clamp_min(1e-12)))
    return torch.stack(losses).mean()


def train(seed: int, k: int, vocab: int, length: int, alpha: float, steps: int, bos: bool, device: torch.device):
    rng = np.random.default_rng(seed)
    torch.manual_seed(seed)
    beta = torch.nn.Parameter(torch.full((k,), 0.25, dtype=torch.float64, device=device))
    log_bos = torch.nn.Parameter(torch.tensor(math.log(vocab), dtype=torch.float64, device=device)) if bos else None
    params = [beta] + ([log_bos] if log_bos is not None else [])
    opt = torch.optim.Adam(params, lr=0.04)
    trace = []
    for step in range(steps):
        seq, _ = sample_chain(k, vocab, length, alpha, rng)
        loss = episode_loss(seq, beta, vocab, log_bos, start=max(8, k + 1))
        opt.zero_grad(); loss.backward(); opt.step()
        if step % 25 == 0 or step == steps - 1:
            trace.append({"step": step, "loss": float(loss.detach()), "beta": beta.detach().tolist(),
                          "bos_mass": float(torch.exp(log_bos.detach())) if log_bos is not None else None})
    return beta.detach(), log_bos.detach() if log_bos is not None else None, trace


def evaluate(seed: int, k: int, vocab: int, length: int, alpha: float, beta: Tensor, log_bos: Tensor | None,
             episodes: int = 40):
    rng = np.random.default_rng(seed + 10000)
    nll = kl_true = kl_bayes = bayes_kl_true = 0.0
    n = 0
    overlap_logw, overlap_degree = [], []
    for _ in range(episodes):
        seq, transition = sample_chain(k, vocab, length, alpha, rng)
        for t in range(max(8, k + 1), length):
            pred = predict(seq, t, beta, vocab, log_bos)
            truth = torch.tensor(transition[context_index(seq[t-k:t], vocab)], dtype=torch.float64, device=beta.device)
            bayes = bayes_predict(seq, t, k, vocab, alpha, beta.device)
            nll += float(-torch.log(pred[seq[t]].clamp_min(1e-12)))
            kl_true += float(torch.sum(truth * torch.log((truth / pred.clamp_min(1e-12)).clamp_min(1e-12))))
            kl_bayes += float(torch.sum(bayes * torch.log((bayes / pred.clamp_min(1e-12)).clamp_min(1e-12))))
            bayes_kl_true += float(torch.sum(truth * torch.log((truth / bayes).clamp_min(1e-12))))
            q = seq[t-k:t]
            for s in range(k, t):
                mask = [int(a == b) for a, b in zip(seq[s-k:s], q)]
                overlap_degree.append(sum(mask))
                overlap_logw.append(sum(float(beta[i]) * mask[i] for i in range(k)))
            n += 1
    x, y = np.asarray(overlap_degree), np.asarray(overlap_logw)
    coef = np.polyfit(x, y, 1)
    fitted = np.polyval(coef, x)
    r2 = 1.0 - float(np.sum((y - fitted) ** 2) / max(np.sum((y - y.mean()) ** 2), 1e-12))
    return {"nll": nll/n, "kl_to_true": kl_true/n, "kl_to_bayes": kl_bayes/n,
            "bayes_kl_to_true": bayes_kl_true/n, "attention_overlap_slope": float(coef[0]),
            "attention_overlap_r2": r2, "positions": n}


def mutation_metrics(seed: int, k: int, vocab: int, length: int, alpha: float, beta: Tensor):
    rng = np.random.default_rng(seed + 20000)
    seq, _ = sample_chain(k, vocab, length, alpha, rng)
    start = max(8, k + 1)
    learned = float(episode_loss(seq, beta, vocab, None, start))
    untrained = float(episode_loss(seq, torch.zeros_like(beta), vocab, None, start))
    hard = float(episode_loss(seq, torch.full_like(beta, 100.0), vocab, None, start))
    shuffled = seq[:start] + random.Random(seed).sample(seq[start:], len(seq[start:]))
    shuffled_loss = float(episode_loss(shuffled, beta, vocab, None, start))
    return {"learned_nll": learned, "untrained_uniform_overlap_nll": untrained,
            "hard_match_nll": hard, "shuffled_context_nll": shuffled_loss,
            "untrained_worse": untrained > learned, "hard_match_worse": hard > learned,
            "shuffled_worse": shuffled_loss > learned}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--output", default="outputs")
    ap.add_argument("--steps", type=int, default=250)
    ap.add_argument("--length", type=int, default=64)
    ap.add_argument("--vocab", type=int, default=5)
    ap.add_argument("--seeds", type=int, default=3)
    ap.add_argument("--device", default="auto", choices=("auto", "cpu", "cuda"))
    ap.add_argument("--orders", type=int, nargs="+", default=[1, 2, 3])
    args = ap.parse_args()
    out = Path(args.output); out.mkdir(parents=True, exist_ok=True)
    device = torch.device("cuda" if args.device == "auto" and torch.cuda.is_available() else args.device if args.device != "auto" else "cpu")
    if device.type == "cuda":
        print(json.dumps({"device": str(device), "gpu": torch.cuda.get_device_name(0)}), flush=True)
    started = time.time(); rows, traces, mutations = [], {}, []
    for k in args.orders:
        for bos in (False, True):
            for seed in range(args.seeds):
                beta, log_bos, trace = train(seed, k, args.vocab, args.length, 1.0, args.steps, bos, device)
                metrics = evaluate(seed, k, args.vocab, args.length, 1.0, beta, log_bos)
                row = {"k": k, "bos": bos, "seed": seed, "beta": beta.tolist(),
                       "bos_mass": float(torch.exp(log_bos)) if log_bos is not None else None, **metrics}
                rows.append(row); traces[f"k{k}_bos{int(bos)}_seed{seed}"] = trace
                if not bos:
                    mutations.append({"k": k, "seed": seed, **mutation_metrics(seed, k, args.vocab, args.length, 1.0, beta)})
                print(json.dumps(row), flush=True)
    flat_rows = [{**r, "beta": json.dumps(r["beta"])} for r in rows]
    with (out / "results.csv").open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=flat_rows[0].keys()); w.writeheader(); w.writerows(flat_rows)
    summary = {
        "paper_setup": {"vocab": args.vocab, "primary_k": 2, "alpha": 1.0},
        "run_setup": vars(args), "wall_seconds": time.time() - started,
        "rows": rows, "mutation_tests": mutations,
        "aggregate": {}
    }
    for k in args.orders:
        for bos in (False, True):
            subset = [r for r in rows if r["k"] == k and r["bos"] == bos]
            summary["aggregate"][f"k{k}_bos{int(bos)}"] = {
                key: {"mean": float(np.mean([r[key] for r in subset])), "std": float(np.std([r[key] for r in subset]))}
                for key in ("nll", "kl_to_true", "kl_to_bayes", "bayes_kl_to_true", "attention_overlap_slope", "attention_overlap_r2")
            }
    (out / "summary.json").write_text(json.dumps(summary, indent=2))
    (out / "training_traces.json").write_text(json.dumps(traces, indent=2))
    print(json.dumps(summary["aggregate"], indent=2))


if __name__ == "__main__":
    main()