Example: add non-gradient attribution baselines

examples/bank_baselines/README.md

@@ -0,0 +1,41 @@
+# Non-gradient attribution baselines
+
+Text similarity baselines evaluated by mean LDS over 50 test queries.
+
+- `bm25_baseline.py` — BM25 lexical overlap: pure surface-form term overlap, no model or embedding.
+- `gradient_baseline.py` — gradient cosine similarity: cosine of the full per-example loss gradients on the bank's model (TracIn-style, unpreconditioned).
+- `activation_baseline.py` — activation similarity: each doc is the mean-pooled input activation to every linear matrix of the model, L2-normalized per matrix and concatenated, then cosine similarity.
+- `semantic_baseline.py` — semantic search with `jinaai/jina-embeddings-v3` (asymmetric `retrieval.query`/`retrieval.passage`).
+- `qwen3_baseline.py` — semantic search with `Qwen/Qwen3-Embedding-8B`, a SOTA decoder embedder (`--model` to swap).
+
+Each produces a `[num_train_docs, num_queries]` score matrix.
+
+## Running
+
+Point a baseline at a re-train bank (written with `save_retrained_models=true`); it reads the model and dataset from the bank's `config.yaml`:
+
+```bash
+python -m examples.bank_baselines.bm25_baseline        --bank runs/retrain_bank_path
+python -m examples.bank_baselines.gradient_baseline    --bank runs/retrain_bank_path
+python -m examples.bank_baselines.activation_baseline  --bank runs/retrain_bank_path
+python -m examples.bank_baselines.semantic_baseline    --bank runs/retrain_bank_path
+python -m examples.bank_baselines.qwen3_baseline       --bank runs/retrain_bank_path
+```
+
+Omit `--bank` to build the default GPT-2/WikiText bank (`examples/magic/gpt2_wikitext_bank.yaml`) first. `--query_split` sets the query set (default `test[1:51]`), `--out` the output dir (default `runs/bank_baselines/`).
+
+## Results
+
+GPT-2 / WikiText bank (100 subsets, 1% leave-out, `eps_root 1e-8`):
+
+| Method | mean ρ |
+| --- | --- |
+| BM25 lexical overlap | 0.16 |
+| Qwen3-Embedding-8B | 0.11 |
+| activation similarity | 0.09 |
+| Jina v3 semantic search | 0.06 |
+| gradient cosine similarity | 0.05 |
+
+## Notes
+
+`jina-embeddings-v3`'s custom code predates transformers 5.x; `semantic_baseline.load_model` sets an `all_tied_weights_keys` default and resets the NaN LoRA `lora_dropout_mask` buffers to ones so it loads and runs. NVIDIA's NV-Embed-v2 is a comparable SOTA embedder but its custom code is incompatible with transformers 5.x, so `qwen3_baseline.py` uses Qwen3-Embedding instead.

examples/bank_baselines/activation_baseline.py

@@ -0,0 +1,153 @@
+"""Baseline: activation similarity over the attribution target modules.
+
+Represents each document by the mean-pooled input activation to every target
+linear module of the bank's fully-trained model (the same modules the gradient
+methods attribute through: all Conv1D/Linear except ``lm_head``). Each
+per-module pooled vector is L2-normalized, then concatenated, so cosine
+similarity between two docs equals the average per-module cosine similarity.
+
+A training doc that is activation-similar to a query is predicted to be
+influential (removing it should raise query loss), so the loss-diff-convention
+score is ``-cosine`` -- larger => larger predicted loss reduction from keeping
+the doc.
+
+Run with (builds the default bank if --bank is omitted):
+    python -m examples.bank_baselines.activation_baseline --bank runs/retrain_bank_path
+"""
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers.pytorch_utils import Conv1D
+
+from . import common
+
+
+def target_module_names(model: torch.nn.Module) -> list[str]:
+    """Linear/Conv1D modules the gradient methods attribute through (not lm_head)."""
+    names = []
+    for name, mod in model.named_modules():
+        if name.endswith("lm_head"):
+            continue
+        if isinstance(mod, (Conv1D, torch.nn.Linear)):
+            names.append(name)
+    return names
+
+
+@torch.no_grad()
+def embed(
+    texts: list[str],
+    model: torch.nn.Module,
+    tokenizer,
+    module_names: list[str],
+    device: str,
+    max_len: int,
+    batch_size: int,
+) -> np.ndarray:
+    """Per-module-normalized, concatenated mean-pooled activation embeddings."""
+    modules = dict(model.named_modules())
+    captured: dict[str, torch.Tensor] = {}
+    handles = []
+    for mn in module_names:
+
+        def hook(_m, inp, _out, _mn=mn):
+            captured[_mn] = inp[0].detach()
+
+        handles.append(modules[mn].register_forward_hook(hook))
+
+    out_rows = []
+    try:
+        for start in range(0, len(texts), batch_size):
+            batch = texts[start : start + batch_size]
+            enc = tokenizer(
+                batch,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=max_len,
+            ).to(device)
+            mask = enc["attention_mask"].unsqueeze(-1).float()  # [B, T, 1]
+            captured.clear()
+            model(**enc)
+
+            per_module = []
+            denom = mask.sum(dim=1).clamp_min(1.0)  # [B, 1]
+            for mn in module_names:
+                act = captured[mn].float()  # [B, T, d]
+                pooled = (act * mask).sum(dim=1) / denom  # [B, d]
+                pooled = torch.nn.functional.normalize(pooled, dim=-1)
+                per_module.append(pooled)
+            emb = torch.cat(per_module, dim=-1)  # [B, sum_d]
+            out_rows.append(emb.cpu().numpy())
+    finally:
+        for h in handles:
+            h.remove()
+    return np.concatenate(out_rows, axis=0)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted")
+    ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT)
+    ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines"))
+    ap.add_argument("--device", default="cuda:0")
+    ap.add_argument("--max_len", type=int, default=1024)
+    ap.add_argument("--batch_size", type=int, default=32)
+    args = ap.parse_args()
+
+    bank = common.ensure_bank(args.bank)
+    spec = common.read_bank_spec(bank)
+    out_dir = Path(args.out)
+
+    tokenizer = AutoTokenizer.from_pretrained(spec.model)
+    tokenizer.pad_token = tokenizer.eos_token
+    model = AutoModelForCausalLM.from_pretrained(
+        spec.base_model, dtype=torch.float32, attn_implementation="eager"
+    ).to(args.device)
+    model.eval()
+
+    module_names = target_module_names(model)
+    print(f"Embedding over {len(module_names)} modules, e.g. {module_names[:3]} ...")
+
+    train_texts, query_texts = common.load_texts(spec, args.query_split)
+    print(f"Embedding {len(train_texts)} train docs ...")
+    train_emb = embed(
+        train_texts,
+        model,
+        tokenizer,
+        module_names,
+        args.device,
+        args.max_len,
+        args.batch_size,
+    )
+    print(f"Embedding {len(query_texts)} query docs ...")
+    query_emb = embed(
+        query_texts,
+        model,
+        tokenizer,
+        module_names,
+        args.device,
+        args.max_len,
+        args.batch_size,
+    )
+
+    # Cosine similarity (rows already unit-norm per module; renormalize the
+    # concatenation so cosine == mean per-module cosine).
+    tn = train_emb / np.linalg.norm(train_emb, axis=1, keepdims=True)
+    qn = query_emb / np.linalg.norm(query_emb, axis=1, keepdims=True)
+    cosine = tn @ qn.T  # [num_train, num_query]
+
+    scores = -cosine  # loss-diff convention (similar => influential)
+    score_path = common.save_scores(scores, out_dir, "activation_scores")
+
+    rhos = common.evaluate_lds(
+        bank, score_path, out_dir / "activation_validate", spec, args.query_split
+    )
+    common.report("activation similarity", rhos)
+
+
+if __name__ == "__main__":
+    main()

examples/bank_baselines/bm25_baseline.py

@@ -0,0 +1,69 @@
+"""Baseline: BM25 lexical overlap.
+
+Scores each training doc against each query by Okapi BM25 -- pure surface-form
+term overlap, no model and no learned embedding. A lexical proxy for influence:
+docs sharing query terms are predicted influential, so the loss-diff-convention
+score is ``-bm25``. Often a stronger influence proxy than deep-semantic
+embedders for small LMs, and essentially free.
+
+Run with (builds the default bank if --bank is omitted):
+    python -m examples.bank_baselines.bm25_baseline --bank runs/retrain_bank_path
+"""
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+from sklearn.feature_extraction.text import CountVectorizer
+
+from . import common
+
+
+def bm25_scores(
+    train_texts: list[str], query_texts: list[str], k1: float = 1.5, b: float = 0.75
+) -> np.ndarray:
+    """Okapi BM25 score of every train doc against every query -> [docs, queries]."""
+    vectorizer = CountVectorizer()
+    tf = vectorizer.fit_transform(train_texts).astype(np.float64).tocsr()  # [D, V]
+    n_docs, _ = tf.shape
+
+    doc_len = np.asarray(tf.sum(axis=1)).ravel()
+    avgdl = doc_len.mean()
+    df = np.asarray((tf > 0).sum(axis=0)).ravel()
+    idf = np.log((n_docs - df + 0.5) / (df + 0.5) + 1.0)
+
+    # Per-nonzero BM25 term weight: idf * tf*(k1+1) / (tf + k1*(1-b+b*dl/avgdl)).
+    rows, cols = tf.nonzero()
+    denom_bias = k1 * (1 - b + b * doc_len[rows] / avgdl)
+    weighted = tf.copy()
+    weighted.data = idf[cols] * tf.data * (k1 + 1.0) / (tf.data + denom_bias)
+
+    # Query terms as a binary term set (Okapi BM25 ignores query term frequency).
+    q_bin = (vectorizer.transform(query_texts) > 0).astype(np.float64)  # [Q, V]
+    return np.asarray((weighted @ q_bin.T).todense())  # [D, Q]
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted")
+    ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT)
+    ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines"))
+    args = ap.parse_args()
+
+    bank = common.ensure_bank(args.bank)
+    spec = common.read_bank_spec(bank)
+    out_dir = Path(args.out)
+
+    train_texts, query_texts = common.load_texts(spec, args.query_split)
+    print(f"BM25 over {len(train_texts)} train docs, {len(query_texts)} queries ...")
+    scores = -bm25_scores(train_texts, query_texts)  # loss-diff convention
+    score_path = common.save_scores(scores, out_dir, "bm25_scores")
+
+    rhos = common.evaluate_lds(
+        bank, score_path, out_dir / "bm25_validate", spec, args.query_split
+    )
+    common.report("BM25 lexical overlap", rhos)
+
+
+if __name__ == "__main__":
+    main()