Add highlighted generation, per-example SHAP, and headline ablation panel

app.py

@@ -1,14 +1,22 @@
 """Predicting Memorization Before Fine-Tuning — interactive demo.
 
 Two modes:
-  (1) Explore the held-out 1M   — pick a real training sequence, see whether the
-      fine-tuned Pythia-1.4B actually memorized it (ground truth), alongside the
-      classifier's pre-fine-tuning risk score and feature breakdown.
+  (1) Explore the held-out 1M   — pick a real training sequence, see the fine-tuned
+      model's actual 50-token generation (green where it matches the ground truth),
+      the memorized/not verdict, and a per-example explanation of the classifier's
+      pre-fine-tuning risk score.
   (2) Score your own text        — paste a passage, get the classifier's predicted
-      memorization risk from base-model features (a forecast; no ground truth).
+      memorization risk from base-model features, with the same per-example explanation.
+
+A collapsible panel shows the headline result: base-model gradient features alone
+predict memorization about as well as the full feature set.
 """
-import os, glob, joblib, numpy as np, pandas as pd, torch, gradio as gr
-import spaces  # Hugging Face ZeroGPU; the @spaces.GPU decorator is a no-op off ZeroGPU
+import os, glob, json, joblib, numpy as np, pandas as pd, torch, gradio as gr
+import spaces  # Hugging Face ZeroGPU; @spaces.GPU is a no-op off ZeroGPU
+import shap
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
 
 HERE = os.path.dirname(os.path.abspath(__file__))
 BASE_MODEL = os.environ.get("BASE_MODEL", "EleutherAI/pythia-1.4b")
@@ -19,35 +27,88 @@ CLFS = {}
 for f in glob.glob(os.path.join(HERE, "models", "clf_*.joblib")):
     b = joblib.load(f)
     CLFS[b["dataset"]] = b
-FEATS = ["zlib_entropy", "base_ppl", "base_loss_variance",
-         "gradient_norm", "one_step_abs_change", "one_step_rel_change"]
-
-# ZeroGPU: place the model on cuda at module level (a CUDA-emulation mode makes this
-# valid at startup; a real GPU is attached only inside @spaces.GPU functions). Off
-# ZeroGPU we fall back to a real GPU if present, else CPU, so the app still runs locally.
+EXPLAINERS = {}  # lazily built shap.TreeExplainer per classifier
+
+FEAT_NAMES = {"zlib_entropy": "zlib entropy", "base_ppl": "base perplexity",
+              "base_loss_variance": "base loss variance", "gradient_norm": "gradient norm",
+              "one_step_abs_change": "one-step loss change", "one_step_rel_change": "one-step rel. loss change"}
+GRAD_FEATS = {"gradient_norm", "one_step_abs_change", "one_step_rel_change"}
+GREEN, RED, GREY = "#2e7d32", "#c62828", "#9aa0a6"
+
+# Headline ablation result (paper: dedup, mean +/- std over 3 training seeds), demo datasets only.
+HEADLINE = {
+    "FineWeb":     {"text-only": (0.628, 0.011), "gradient-only": (0.994, 0.003), "all": (0.997, 0.002)},
+    "PG-19":       {"text-only": (0.640, 0.035), "gradient-only": (0.998, 0.001), "all": (0.996, 0.003)},
+    "The Stack":   {"text-only": (0.765, 0.001), "gradient-only": (0.965, 0.001), "all": (0.971, 0.001)},
+    "OpenWebMath": {"text-only": (0.837, 0.001), "gradient-only": (0.976, 0.002), "all": (0.992, 0.001)},
+}
+
+# ZeroGPU: place the base model on cuda at module level (a CUDA-emulation mode makes
+# this valid at startup; a real GPU is attached only inside @spaces.GPU functions).
 _ON_ZEROGPU = bool(os.environ.get("SPACES_ZERO_GPU"))
 DEVICE = "cuda" if (_ON_ZEROGPU or torch.cuda.is_available()) else "cpu"
 try:
     from byo_features import BYOExtractor
     EXTRACTOR = BYOExtractor(BASE_MODEL, device=DEVICE)
     _LOAD_ERR = None
-except Exception as e:  # e.g. no model cached / offline local run
+except Exception as e:
     EXTRACTOR, _LOAD_ERR = None, str(e)
 
 
 @spaces.GPU(duration=120)
 def _compute_features(text):
-    """GPU-bound step of the BYO path (forward + backward on the base model)."""
     return EXTRACTOR.features(text)
 
 
-def _feature_table(row):
-    names = {"zlib_entropy": "zlib entropy", "base_ppl": "base perplexity",
-             "base_loss_variance": "base loss variance", "gradient_norm": "gradient norm",
-             "one_step_rel_change": "one-step rel. loss change"}
-    return pd.DataFrame(
-        [(v, f"{float(row[k]):.4f}") for k, v in names.items() if k in row],
-        columns=["Base-model feature", "Value"])
+# ------------------------------- graphics ----------------------------------
+def headline_figure():
+    tiers = ["text-only", "gradient-only", "all"]
+    colors = {"text-only": GREY, "gradient-only": GREEN, "all": "#1f4e79"}
+    datasets = list(HEADLINE.keys())
+    x = np.arange(len(datasets)); w = 0.26
+    fig, ax = plt.subplots(figsize=(7.2, 3.6))
+    for i, t in enumerate(tiers):
+        vals = [HEADLINE[d][t][0] for d in datasets]
+        errs = [HEADLINE[d][t][1] for d in datasets]
+        ax.bar(x + (i - 1) * w, vals, w, yerr=errs, capsize=3, label=t, color=colors[t])
+    ax.set_xticks(x); ax.set_xticklabels(datasets)
+    ax.set_ylim(0.5, 1.02); ax.set_ylabel("ROC-AUC")
+    ax.set_title("Gradient-only features ≈ full model (mean ± std over 3 seeds)")
+    ax.axhline(0.5, color="k", lw=0.7, ls=":")
+    ax.legend(loc="lower right", ncol=3, fontsize=8, framealpha=0.9)
+    fig.tight_layout()
+    return fig
+
+
+def shap_figure(feats, clf_name):
+    b = CLFS[clf_name]
+    if clf_name not in EXPLAINERS:
+        EXPLAINERS[clf_name] = shap.TreeExplainer(b["clf"])
+    expl = EXPLAINERS[clf_name]
+    order = b["feats"]
+    X = np.nan_to_num(np.array([[feats[k] for k in order]], dtype=float), posinf=1e10, neginf=-1e10)
+    Xs = b["scaler"].transform(X)
+    sv = expl.shap_values(Xs)
+    sv = np.asarray(sv[1] if isinstance(sv, list) else sv).reshape(-1)[:len(order)]
+
+    idx = np.argsort(np.abs(sv))               # ascending -> largest on top
+    labels = [f"{FEAT_NAMES.get(order[i], order[i])}  ({feats[order[i]]:.3g})" for i in idx]
+    vals = sv[idx]
+    colors = [GREEN if v > 0 else GREY for v in vals]
+    fig, ax = plt.subplots(figsize=(7.2, 3.2))
+    ax.barh(range(len(vals)), vals, color=colors)
+    ax.set_yticks(range(len(vals))); ax.set_yticklabels(labels, fontsize=9)
+    ax.axvline(0, color="k", lw=0.8)
+    ax.set_xlabel("push toward memorized  →   (log-odds contribution)")
+    grad_top = [order[i] in GRAD_FEATS for i in idx]
+    ax.set_title("Why this score" + ("  ·  gradient features dominate" if any(grad_top[-2:]) else ""))
+    fig.tight_layout()
+    return fig
+
+
+def _segments(gen_segments_json):
+    segs = json.loads(gen_segments_json)
+    return [(t, lab) for t, lab in segs]
 
 
 # ----------------------------- Explore (lookup) -----------------------------
@@ -58,18 +119,23 @@ def draw(dataset, which):
     elif which == "Not memorized":
         df = df[~df.memorized]
     if len(df) == 0:
-        return "No examples for this filter.", "", "", None
+        return "No examples for this filter.", [], "", None
     r = df.sample(1).iloc[0]
+    gen = _segments(r.gen_segments)
     if r.memorized:
-        verdict = "### 🔴 Memorized\nThe fine-tuned model reproduced the ground-truth suffix exactly."
+        verdict = "### 🔴 Memorized\nThe fine-tuned model reproduced the ground-truth suffix exactly (all tokens green)."
     else:
-        verdict = "### 🟢 Not memorized\nThe fine-tuned model did not reproduce the ground-truth suffix."
+        verdict = ("### 🟢 Not memorized\nThe fine-tuned model's greedy generation diverges from the ground truth; "
+                   "green tokens happen to match, red tokens do not.")
     if bool(r.base_reproduced):
         verdict += "\n\n_(Note: the base model already reproduces this suffix, so it is contamination rather than fine-tuning-induced memorization.)_"
-    score_md = (f"### Classifier risk score: **{float(r.score):.3f}**\n"
-                f"Computed from base-model features *before* fine-tuning. "
-                f"Higher means more likely to be memorized. Ground truth above is what actually happened.")
-    return (r.prefix, r.suffix, verdict + "\n\n" + score_md, _feature_table(r))
+    verdict += (f"\n\n### Classifier risk score: **{float(r.score):.3f}**\n"
+                f"Computed from base-model features *before* fine-tuning. The chart below shows how each "
+                f"feature moved this score.")
+    clf_name = dataset if dataset in CLFS else next(iter(CLFS))
+    feats = {k: float(r[k]) for k in CLFS[clf_name]["feats"]}
+    fig = shap_figure(feats, clf_name)
+    return r.prefix, gen, verdict, fig
 
 
 # ----------------------------- Predict (BYO) --------------------------------
@@ -89,9 +155,9 @@ def predict(text, clf_name):
     flag = "high" if risk >= 0.5 else "low"
     md = (f"### Predicted memorization risk: **{risk:.3f}** ({flag})\n"
           f"Forecast from base-model features using the classifier trained on **{clf_name}**. "
-          f"This is a prediction of whether a model *would* memorize this sequence if fine-tuned on it; "
-          f"there is no ground truth for arbitrary text. Treat it as a relative risk score, not a calibrated probability.")
-    return md, _feature_table(feats)
+          f"There is no ground truth for arbitrary text, so treat it as a relative risk score. "
+          f"The chart shows which features drove it.")
+    return md, shap_figure(feats, clf_name)
 
 
 INTRO = """# Predicting Memorization Before Fine-Tuning
@@ -100,30 +166,42 @@ fine-tuned model will memorize. **Explore** verifies this against ground truth o
 run; **Score your own text** applies the same classifier to a passage you provide.
 """
 
+WHY = """The classifier reads six features of the **base** model, measured **before** any fine-tuning.
+Three come from a single gradient step on the sequence (**gradient norm** and the **one-step loss change**);
+the rest are text and base-loss statistics. On the datasets shown here, the three **gradient-only** features
+alone reach ROC-AUC 0.965–0.998, essentially matching the full feature set, while text-only features lag far
+behind. (On conversational text, not shown here, the non-gradient features contribute more.)
+"""
+
 with gr.Blocks(title="Predicting Memorization Before Fine-Tuning") as demo:
     gr.Markdown(INTRO)
     with gr.Tab("Explore the held-out 1M"):
+        with gr.Accordion("Why base-model gradients predict memorization", open=False):
+            gr.Markdown(WHY)
+            gr.Plot(headline_figure())
         gr.Markdown("Pick a real Pythia-1.4B training sequence and compare the classifier's "
-                    "pre-fine-tuning risk score with whether the fine-tuned model actually memorized it.")
+                    "pre-fine-tuning risk score with what the fine-tuned model actually generated.")
         with gr.Row():
             d_in = gr.Dropdown(DATASETS, value=DATASETS[0], label="Dataset")
             f_in = gr.Radio(["All", "Memorized only", "Not memorized"], value="All", label="Filter")
             go = gr.Button("Show a random example", variant="primary")
         pref = gr.Textbox(label="Prefix (50 tokens given to the model)", lines=3)
-        suf = gr.Textbox(label="Ground-truth suffix (next 50 tokens)", lines=3)
+        gen = gr.HighlightedText(
+            label="Model's actual 50-token generation (green = matches ground truth, red = differs)",
+            color_map={"match": GREEN, "diff": RED}, combine_adjacent=True, show_legend=False)
         verdict = gr.Markdown()
-        feat = gr.Dataframe(label="Base-model features for this sequence", interactive=False)
-        go.click(draw, [d_in, f_in], [pref, suf, verdict, feat])
+        feat = gr.Plot(label="Per-example feature contributions (SHAP)")
+        go.click(draw, [d_in, f_in], [pref, gen, verdict, feat])
     with gr.Tab("Score your own text"):
         gr.Markdown("Paste a passage (at least ~100 tokens). We compute its base-model features and "
-                    "return the classifier's **predicted** memorization risk. No ground truth exists for "
-                    "arbitrary text, so this is a forecast (see the trust discussion in the paper).")
+                    "return the classifier's **predicted** memorization risk, with the per-example "
+                    "feature contributions below.")
         txt = gr.Textbox(label="Your text", lines=6)
         clf_in = gr.Dropdown(sorted(CLFS.keys()), value=sorted(CLFS.keys())[0],
                              label="Classifier (trained on)")
         pgo = gr.Button("Predict memorization risk", variant="primary")
         pmd = gr.Markdown()
-        pfeat = gr.Dataframe(label="Base-model features for your text", interactive=False)
+        pfeat = gr.Plot(label="Per-example feature contributions (SHAP)")
         pgo.click(predict, [txt, clf_in], [pmd, pfeat])
     gr.Markdown("_Datasets shown are public (FineWeb, PG-19, The Stack, OpenWebMath). "
                 "The classifier uses only base-model features and was trained on a separate fine-tuning run._")

data/lookup.parquet

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4bcf4588bd84e05e61578ff3d7794f505ec2917c3a88f2bedfe93cd1094ab114
-size 6934910
+oid sha256:2d4a15810da1cc066350a3aea465eef483522cdfbd3b4d70361f5eeb5a2e6057
+size 11963509

requirements.txt

@@ -7,3 +7,5 @@ pandas>=2.0
 pyarrow>=14.0
 numpy>=2.0
 joblib>=1.5
+shap==0.49.1
+matplotlib>=3.7