bundle/scripts/claim1_decomposition.py

#!/usr/bin/env python
"""
Claim 1: MFA decomposes activations into a discrete region assignment (responsibilities
over centroids) plus a continuous within-region offset (posterior-mean loadings),
formalized in Eqs 5-13.

We verify the released MFA code implements Eqs 5-13 exactly and that the decomposition
x ~= A b(x)  (Eq 11-13) reconstructs real Gemma-2-2B residual-stream activations.

Checks:
  (a) Responsibilities R_k(x) (Eq 8) are a valid posterior: nonneg, sum to 1.
  (b) Posterior-mean latent z_hat_k (Eqs 9-10) equals the closed-form Z_k (x - mu_k).
  (c) Dictionary A = [mu_1 | W_1 | ... | mu_K | W_K] (Eq 12) and coefficients
      b(x) = [R_k, R_k z_hat_k] (Eq 13) give recon = A b(x) (Eq 11).
  (d) Two-segment structure: centroid term explains region-level, offset refines it.
      Report relative reconstruction error and R^2 vs. the model's own reconstruct().
"""
import os, sys, json, time, argparse
import torch
import numpy as np

sys.path.insert(0, os.environ.get("MFA_REPO", os.path.join(os.path.dirname(__file__), "..", "repo")))
from modeling.model_checkpointing import load_mfa
from modeling.mfa import MFAEncoderDecoder


def get_activations(model_name, layer, n_tokens, device, seq_len=64):
    """Extract residual-stream activations (resid_post) at `layer` from real text."""
    from transformer_lens import HookedTransformer
    from datasets import load_dataset
    tl = HookedTransformer.from_pretrained(model_name, device=device, dtype=torch.float32)
    hook = f"blocks.{layer}.hook_resid_post"
    ds = load_dataset("NeelNanda/pile-10k", split="train", streaming=True)
    acts = []
    for ex in ds:
        toks = tl.to_tokens(ex["text"])[:, :seq_len]
        with torch.no_grad():
            _, cache = tl.run_with_cache(toks, names_filter=hook, return_type=None)
        a = cache[hook].reshape(-1, tl.cfg.d_model).float().cpu()
        # drop BOS position per example
        acts.append(a[1:])
        if sum(x.shape[0] for x in acts) >= n_tokens:
            break
    X = torch.cat(acts, 0)[:n_tokens]
    del tl
    torch.cuda.empty_cache()
    return X


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--model", default="gemma-2-2b")
    ap.add_argument("--layer", type=int, default=18)
    ap.add_argument("--n", type=int, default=4096)
    ap.add_argument("--out", default="outputs/claim1_decomposition.json")
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    t0 = time.time()

    mfa = load_mfa(args.ckpt, map_location=device).eval()
    print(f"Loaded MFA: K={mfa.K} D={mfa.D} q={mfa.q}")

    cache = f"outputs/acts_{args.model}_l{args.layer}_{args.n}.pt"
    if os.path.exists(cache):
        X = torch.load(cache).to(device)
        print(f"Loaded cached activations {tuple(X.shape)}")
    else:
        X = get_activations(args.model, args.layer, args.n, device).to(device)
        os.makedirs("outputs", exist_ok=True)
        torch.save(X.cpu(), cache)
    print(f"Activations: {tuple(X.shape)} from {args.model} L{args.layer}")

    with torch.no_grad():
        # (a) responsibilities
        R = mfa.responsibilities(X)                     # (B,K)  Eq 8
        resp_sum = R.sum(1)
        resp_min = R.min().item()
        assign = R.argmax(1)
        used = assign.unique().numel()

        # (b) posterior mean z_hat via model, and closed-form Z_k(x-mu_k) (Eqs 9-10)
        Ez, _ = mfa.component_posterior(X)              # (B,K,q)
        # closed form for each token's assigned component
        W = mfa.W                                       # (K,D,q)
        psi = mfa._psi()                                # (K,D)
        idx = assign
        Wk = W[idx]                                     # (B,D,q)
        psik = psi[idx]                                 # (B,D)
        muk = mfa.mu[idx]                               # (B,D)
        Iq = torch.eye(mfa.q, device=device)
        # Z_k = (I + W^T Psi^-1 W)^-1 W^T Psi^-1
        WtPinv = torch.einsum("bdq,bd->bqd", Wk, 1.0/psik)   # (B,q,D)
        M = Iq[None] + torch.einsum("bqd,bdr->bqr", WtPinv, Wk)  # (B,q,q)
        zclosed = torch.linalg.solve(M, torch.einsum("bqd,bd->bq", WtPinv, (X-muk)))
        zmodel = Ez[torch.arange(X.shape[0]), idx]      # (B,q)
        z_err = (zclosed - zmodel).norm() / zmodel.norm().clamp_min(1e-9)

        # (c) full dictionary decomposition Eqs 11-13, done in token-chunks to bound memory.
        # encode() assembles A = [mu_1|W_1|...|mu_K|W_K] (Eq 12) and
        # b(x) = [R_k, R_k z_k] (Eq 13) then recon = A b(x) = coeffs @ A^T (Eq 11).
        ed = MFAEncoderDecoder(mfa)
        recon_parts = []
        for s in range(0, X.shape[0], 512):
            recon_parts.append(ed.encode(X[s:s+512]).recon)
        recon = torch.cat(recon_parts, 0)               # (B,D) = A b(x)
        # relative reconstruction error of the Eq 11-13 dictionary reconstruction
        rel_err_mix = ((recon - X).norm(dim=1) / X.norm(dim=1)).mean().item()

        # single-component (hard-assigned) two-segment reconstruction:
        # segment 1: centroid only ;  segment 1+2: centroid + W z
        cent = muk                                      # region term  R_k mu_k (hard)
        offset = torch.einsum("bdq,bq->bd", Wk, zmodel) # within-region offset
        err_centroid = ((cent - X).norm(dim=1)/X.norm(dim=1)).mean().item()
        err_two = ((cent + offset - X).norm(dim=1)/X.norm(dim=1)).mean().item()

        # variance explained by hard two-segment recon
        ss_res = ((cent+offset - X)**2).sum().item()
        ss_tot = ((X - X.mean(0))**2).sum().item()
        r2_two = 1 - ss_res/ss_tot

        # dictionary sparsity: nonzero coeff blocks per token (should be ~all K responsibilities
        # but effective support is small -> measure mass in top block)
        Rsort = R.sort(1, descending=True).values
        top1_mass = Rsort[:, 0].mean().item()
        top5_mass = Rsort[:, :5].sum(1).mean().item()

    out = {
        "ckpt": os.path.basename(args.ckpt),
        "model": args.model, "layer": args.layer,
        "K": mfa.K, "D": mfa.D, "q": mfa.q,
        "n_tokens": int(X.shape[0]),
        "resp_sum_mean": float(resp_sum.mean()),
        "resp_sum_max_dev": float((resp_sum-1).abs().max()),
        "resp_min": float(resp_min),
        "components_used": int(used),
        "z_closedform_rel_err": float(z_err),        # (b) should be ~0
        "recon_rel_err_mixture": rel_err_mix,        # (c)
        "err_centroid_only": err_centroid,           # (d) segment 1
        "err_centroid_plus_offset": err_two,         # (d) segment 1+2
        "r2_two_segment": r2_two,
        "top1_responsibility_mass": top1_mass,
        "top5_responsibility_mass": top5_mass,
        "wall_s": time.time()-t0,
    }
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    json.dump(out, open(args.out, "w"), indent=2)
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()