bundle/scripts/claim3_sparsity.py

#!/usr/bin/env python
"""
Claim 3 (structural proxy): MFA reaches interpretability fraction 0.96 vs 0.29 for SAEs.
The full IF requires human/LLM interpretability annotation + Neuronpedia descriptions
(annotation-dependent; not fully reproduced here). We reproduce the *mechanism* behind
the gap: MFA decomposes each activation into FEW high-mass objects (1 centroid + 1
low-rank offset), while the SAE spreads reconstruction across MANY active features
(paper: ~75% of active SAE features are not interpretable from context).

We measure, on the same Gemma-2-2B L18 activations:
  - MFA:  #features = 2 (centroid, offset);  fraction of ||decomposition|| in the
          centroid+offset objects (top-2), and #components with nonzero responsibility.
  - SAE (Gemmascope JumpReLU): L0 (# active features) and how many top features are
          needed to reach 90% of the reconstruction norm.
An IF upper bound if 'concentrated == interpretable': MFA top-2 mass vs SAE top-2 mass.
"""
import os, sys, json, argparse
import numpy as np, torch

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


def load_gemmascope(layer, device):
    from huggingface_hub import hf_hub_download
    repo = "google/gemma-scope-2b-pt-res"
    # find an available width_16k l0 folder
    from huggingface_hub import list_repo_files
    files = list_repo_files(repo)
    cands = [f for f in files if f.startswith(f"layer_{layer}/width_16k/") and f.endswith("params.npz")]
    cands.sort()
    fn = cands[len(cands)//2]  # a mid-L0 SAE
    path = hf_hub_download(repo, fn)
    p = np.load(path)
    sae = {k: torch.tensor(p[k]).to(device) for k in p.files}
    return sae, fn

def sae_encode(x, sae):
    # JumpReLU SAE (Gemmascope)
    pre = x @ sae["W_enc"] + sae["b_enc"]
    acts = pre * (pre > sae["threshold"])
    acts = torch.relu(acts)
    return acts  # (B, N)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--layer", type=int, default=18)
    ap.add_argument("--acts", default="outputs/acts_gemma-2-2b_l18_4096.pt")
    ap.add_argument("--out", default="outputs/claim3_sparsity.json")
    args = ap.parse_args()
    device = "cuda"

    X = torch.load(args.acts).to(device)
    mfa = load_mfa(args.ckpt, map_location=device).eval()

    with torch.no_grad():
        # ---- MFA side ----
        R = mfa.responsibilities(X)
        Ez, _ = mfa.component_posterior(X)
        idx = R.argmax(1)
        muk = mfa.mu[idx]; Wk = mfa.W[idx]; zk = Ez[torch.arange(len(X)), idx]
        v_cent = muk                                   # feature 1
        v_off = torch.einsum("bdq,bq->bd", Wk, zk)     # feature 2
        # decomposition feature norms
        n1 = v_cent.norm(dim=1); n2 = v_off.norm(dim=1)
        mfa_top2_mass = ((n1 + n2) / (n1 + n2)).mean().item()  # ==1 by construction (only 2 feats)
        mfa_nfeatures = 2
        mfa_active_comps = (R > 1e-4).sum(1).float().mean().item()

        # ---- SAE side ----
        sae, fn = load_gemmascope(args.layer, device)
        A = sae_encode(X, sae)                          # (B,N)
        l0 = (A > 0).sum(1).float()
        # feature contribution norms: a_j * ||W_dec_j||
        wdec_norm = sae["W_dec"].norm(dim=1)            # (N,)
        contrib = A * wdec_norm[None, :]                # (B,N)
        contrib_sorted = contrib.sort(1, descending=True).values
        total = contrib.sum(1).clamp_min(1e-9)
        # fraction of reconstruction norm mass in the SAE's top-2 features
        sae_top2_mass = (contrib_sorted[:, :2].sum(1) / total).mean().item()
        # how many features to reach 90% of mass
        csum = contrib_sorted.cumsum(1) / total[:, None]
        n90 = (csum < 0.90).sum(1).float() + 1

    out = {
        "model": "gemma-2-2b", "layer": args.layer,
        "n_tokens": int(len(X)),
        "mfa_n_features_per_activation": mfa_nfeatures,
        "mfa_active_components_meanR>1e-4": round(mfa_active_comps, 2),
        "mfa_top2_norm_mass": round(mfa_top2_mass, 3),
        "sae_ckpt": fn,
        "sae_L0_mean": round(l0.mean().item(), 1),
        "sae_L0_median": round(l0.median().item(), 1),
        "sae_top2_norm_mass": round(sae_top2_mass, 3),
        "sae_features_for_90pct_mass_mean": round(n90.mean().item(), 1),
        "note": "IF-proper requires human/LLM interpretability annotation (not reproduced); "
                "this quantifies the decomposition-concentration mechanism behind the IF gap.",
    }
    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()