From adecdf8fee3cfafb4b68fca6c0ee0798c8dfe330 Mon Sep 17 00:00:00 2001 From: Dang Nguyen Date: Fri, 5 Jun 2026 16:53:30 -0500 Subject: [PATCH] Add generator selection-bias analysis Compares the feature distribution of perturbation spans the generator selected against the full candidate pool it chose from. Random selection would make the two distributions match, so divergence measures selection bias directly, with no new LLM calls. Walks existing *_perturbations.json files and re-extracts the candidate pool with the same extractor. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../perturbation/analyses/selection_bias.py | 386 ++++++++++++++++++ 1 file changed, 386 insertions(+) create mode 100644 benchmarks/perturbation/analyses/selection_bias.py diff --git a/benchmarks/perturbation/analyses/selection_bias.py b/benchmarks/perturbation/analyses/selection_bias.py new file mode 100644 index 0000000..b4c270e --- /dev/null +++ b/benchmarks/perturbation/analyses/selection_bias.py @@ -0,0 +1,386 @@ +"""Quantify generator selection bias: selected perturbation spans vs. the candidate pool. + +The generator LLM picks which candidate spans to perturb (in batches of 10). +Random selection would yield a selected subset whose feature distribution +matches the pool, so any divergence between selected and pool measures +selection bias directly, with no new LLM calls. + +Walks `--input-root` for `*_perturbations.json` files, re-extracts the +candidate pool from the matching `_clean.md` (extraction is deterministic, +so span_ids are reproduced exactly), reconstructs the batches the generator +saw, and marks which candidates were selected (survived generation + +structural validation) and which were kept (survived the substantive-error +verifier, from `--filtered-root`). + +Analyses: + 1. Sanity check: manifest `original` matches the re-extracted span text. + 2. Selection rates per run; runs with 100% selection have no freedom and + are excluded from bias tests. + 3. Span type: selection rate per span_type within mixed-type pools. + 4. Span length: selected vs unselected, within-paper. + 5. Within-batch position: permutation test for primacy bias. + 6. Document position: selection rate by offset quartile. + 7. Pool truncation: how much the 20-per-type extraction cap discards. + 8. Equation error-subtype choice (free choice, no pool baseline). + +Confidence intervals use a cluster bootstrap over papers. + +Usage: + python -m benchmarks.perturbation.analyses.selection_bias + python -m benchmarks.perturbation.analyses.selection_bias --error-type surface +""" + +from __future__ import annotations + +import argparse +import csv +import json +from collections import Counter, defaultdict +from pathlib import Path + +import numpy as np + +from .. import extract as extract_mod +from ..extract import extract_candidates +from ..generate import _BATCH_SIZE +from .verify_existing import _paper_category_for + +REPO_ROOT = Path(__file__).resolve().parents[3] +DEFAULT_INPUT_ROOT = REPO_ROOT / "benchmarks" / "perturbation" / "data" / "perturbations" +DEFAULT_FILTERED_ROOT = REPO_ROOT / "benchmarks" / "perturbation" / "data" / "perturbations_filtered" +DEFAULT_OUTPUT_DIR = REPO_ROOT / "benchmarks" / "perturbation" / "results" / "selection_bias" + +RNG = np.random.default_rng(0) +N_BOOT = 2000 +N_PERM = 10000 + +EQUATION_TYPES = {"equation_display", "equation_inline", "equation_named"} + + +def _collect_rows(input_root: Path, filtered_root: Path, error_type_filter: str | None): + """One row per (run, candidate). Returns (rows, run_stats, sanity).""" + rows: list[dict] = [] + run_stats: list[dict] = [] + sanity = {"n_perturbations": 0, "span_id_missing": 0, "text_mismatch": 0, "runs_skipped": []} + + for perts_path in sorted(input_root.rglob("*_perturbations.json")): + rel = perts_path.relative_to(input_root) + error_type_run = perts_path.parent.name + if error_type_filter and error_type_run != error_type_filter: + continue + slug = perts_path.name.replace("_perturbations.json", "") + clean_md = perts_path.with_name(f"{slug}_clean.md") + if not clean_md.exists(): + sanity["runs_skipped"].append(f"{rel}: no _clean.md") + continue + + manifest = json.loads(perts_path.read_text()) + field = rel.parts[0] + paper = rel.parts[2] # /all/// + text = clean_md.read_text() + category = _paper_category_for(error_type_run) + candidates = extract_candidates(category, error_type_run, text) + + # Sanity: re-extraction must reproduce the pool the generator saw. + if len(candidates) != manifest["n_candidates"]: + sanity["runs_skipped"].append( + f"{rel}: re-extracted {len(candidates)} candidates, manifest says {manifest['n_candidates']}" + ) + continue + + by_id = {c.span_id: c for c in candidates} + selected: dict[str, str] = {} # span_id -> chosen error subtype + for p in manifest["perturbations"]: + sanity["n_perturbations"] += 1 + cand = by_id.get(p["span_id"]) + if cand is None: + sanity["span_id_missing"] += 1 + continue + if p["original"] != cand.text and p["original"] not in cand.text: + sanity["text_mismatch"] += 1 + selected[p["span_id"]] = p["error"] + + kept_ids: set[str] = set() + kept_path = filtered_root / rel.parent / f"{slug}_kept_perturbations.json" + if kept_path.exists(): + kept = json.loads(kept_path.read_text()) + kept_ids = {p["span_id"] for p in kept.get("perturbations", [])} + + # Full pool without the per-type cap, to measure truncation. + cap = extract_mod._MAX_PER_SPAN_TYPE + extract_mod._MAX_PER_SPAN_TYPE = 10 ** 9 + try: + uncapped = extract_candidates(category, error_type_run, text) + finally: + extract_mod._MAX_PER_SPAN_TYPE = cap + uncapped_counts = Counter(c.span_type.value for c in uncapped) + capped_counts = Counter(c.span_type.value for c in candidates) + + batches = [candidates[i:i + _BATCH_SIZE] for i in range(0, len(candidates), _BATCH_SIZE)] + for batch_idx, batch in enumerate(batches): + for pos, c in enumerate(batch): + rows.append({ + "field": field, + "paper": paper, + "error_type_run": error_type_run, + "slug": slug, + "span_id": c.span_id, + "span_type": c.span_type.value, + "len_chars": len(c.text), + "offset_frac": c.offset / max(1, len(text)), + "batch_idx": batch_idx, + "n_batches": len(batches), + "pos_in_batch": pos, + "batch_size": len(batch), + "selected": int(c.span_id in selected), + "kept": int(c.span_id in kept_ids), + "chosen_error": selected.get(c.span_id, ""), + }) + + run_stats.append({ + "run": str(rel.parent), + "error_type_run": error_type_run, + "paper": paper, + "n_candidates": len(candidates), + "n_selected": len(selected), + "n_kept": len(kept_ids), + "n_generated": manifest.get("n_generated"), + "selection_rate": len(selected) / max(1, len(candidates)), + "truncation": { + t: {"uncapped": uncapped_counts[t], "capped": capped_counts[t]} + for t in uncapped_counts + }, + }) + + return rows, run_stats, sanity + + +# --------------------------------------------------------------------------- +# Statistics helpers +# --------------------------------------------------------------------------- + +def _cluster_bootstrap_rate_diff(rows, key_fn, group_a, group_b, outcome="selected"): + """Bootstrap CI (over papers) for rate(group_a) - rate(group_b).""" + by_paper: dict[str, list[dict]] = defaultdict(list) + for r in rows: + by_paper[r["paper"]].append(r) + papers = sorted(by_paper) + + def point(sample_papers): + a_n = a_d = b_n = b_d = 0 + for p in sample_papers: + for r in by_paper[p]: + g = key_fn(r) + if g == group_a: + a_n += r[outcome]; a_d += 1 + elif g == group_b: + b_n += r[outcome]; b_d += 1 + if a_d == 0 or b_d == 0: + return None + return a_n / a_d - b_n / b_d + + est = point(papers) + if est is None: + return None + boots = [] + for _ in range(N_BOOT): + sample = RNG.choice(papers, size=len(papers), replace=True) + v = point(list(sample)) + if v is not None: + boots.append(v) + lo, hi = np.percentile(boots, [2.5, 97.5]) + return {"diff": est, "ci_lo": float(lo), "ci_hi": float(hi)} + + +def _position_permutation_test(rows): + """Within-batch primacy test. + + Statistic: mean within-batch position of selected candidates minus its + expectation under random within-batch selection (averaged over batches + with partial selection). Negative = primacy bias (early items favored). + The null is built by permuting the selected mask within each batch. + """ + batches: dict[tuple, list[dict]] = defaultdict(list) + for r in rows: + batches[(r["slug"], r["error_type_run"], r["paper"], r["batch_idx"])].append(r) + + obs_terms, perm_batches = [], [] + for key, items in batches.items(): + items.sort(key=lambda r: r["pos_in_batch"]) + sel = np.array([r["selected"] for r in items]) + k, m = int(sel.sum()), len(sel) + if k == 0 or k == m or m < 2: + continue # no selection freedom in this batch + pos = np.arange(m) + expected = pos.mean() + obs_terms.append(pos[sel == 1].mean() - expected) + perm_batches.append((pos, k, expected)) + + if not obs_terms: + return None + observed = float(np.mean(obs_terms)) + + null = np.empty(N_PERM) + for i in range(N_PERM): + terms = [ + RNG.choice(pos, size=k, replace=False).mean() - expected + for pos, k, expected in perm_batches + ] + null[i] = np.mean(terms) + p_value = float(np.mean(np.abs(null) >= abs(observed))) + return { + "observed_mean_position_shift": observed, + "p_value": p_value, + "n_partial_batches": len(perm_batches), + "null_sd": float(null.std()), + } + + +def _rate_by_bin(rows, value_fn, n_bins=4, outcome="selected"): + """Selection rate by within-run quantile bin of a numeric feature.""" + by_run: dict[tuple, list[dict]] = defaultdict(list) + for r in rows: + by_run[(r["slug"], r["error_type_run"], r["paper"])].append(r) + + bin_counts = np.zeros(n_bins) + bin_sel = np.zeros(n_bins) + for items in by_run.values(): + vals = np.array([value_fn(r) for r in items], dtype=float) + if len(set(vals)) < 2: + continue + edges = np.quantile(vals, np.linspace(0, 1, n_bins + 1)) + bins = np.clip(np.searchsorted(edges, vals, side="right") - 1, 0, n_bins - 1) + for b, r in zip(bins, items): + bin_counts[b] += 1 + bin_sel[b] += r[outcome] + return { + f"q{i + 1}": {"rate": float(bin_sel[i] / bin_counts[i]), "n": int(bin_counts[i])} + for i in range(n_bins) if bin_counts[i] > 0 + } + + +# --------------------------------------------------------------------------- +# Report +# --------------------------------------------------------------------------- + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--input-root", type=Path, default=DEFAULT_INPUT_ROOT) + ap.add_argument("--filtered-root", type=Path, default=DEFAULT_FILTERED_ROOT) + ap.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR) + ap.add_argument("--error-type", default=None, + help="Restrict to one error_type run dir (e.g. surface)") + args = ap.parse_args() + + rows, run_stats, sanity = _collect_rows(args.input_root, args.filtered_root, args.error_type) + print(f"{len(run_stats)} runs, {len(rows)} candidate rows") + print(f"Sanity: {sanity['n_perturbations']} perturbations, " + f"{sanity['span_id_missing']} missing span_ids, " + f"{sanity['text_mismatch']} text mismatches, " + f"{len(sanity['runs_skipped'])} runs skipped") + for s in sanity["runs_skipped"]: + print(f" SKIPPED {s}") + + # Restrict bias analyses to runs with selection freedom. + full_runs = {s["run"] for s in run_stats if s["selection_rate"] >= 1.0} + free_rows = [ + r for r in rows + if f"{Path(r['field']) / 'all' / r['paper'] / r['error_type_run']}" not in full_runs + ] + print(f"{len(full_runs)} runs selected 100% of candidates (excluded from bias tests); " + f"{len(free_rows)} rows remain") + + summary: dict = { + "n_runs": len(run_stats), + "n_rows": len(rows), + "sanity": sanity, + "n_full_selection_runs": len(full_runs), + "by_error_type": {}, + } + + error_types = sorted({r["error_type_run"] for r in free_rows}) + for et in error_types: + sub = [r for r in free_rows if r["error_type_run"] == et] + et_summary: dict = { + "n_rows": len(sub), + "overall_selection_rate": float(np.mean([r["selected"] for r in sub])), + } + + # 3. Span type rates (only informative for mixed-type pools). + type_rates = {} + for st in sorted({r["span_type"] for r in sub}): + st_rows = [r for r in sub if r["span_type"] == st] + type_rates[st] = { + "rate": float(np.mean([r["selected"] for r in st_rows])), + "n_pool": len(st_rows), + "n_selected": int(sum(r["selected"] for r in st_rows)), + } + et_summary["selection_rate_by_span_type"] = type_rates + types = sorted(type_rates, key=lambda t: -type_rates[t]["n_pool"]) + if len(types) >= 2: + et_summary["span_type_top2_diff"] = _cluster_bootstrap_rate_diff( + sub, lambda r: r["span_type"], types[0], types[1]) + + # 4. Length and 6. document position, by within-run quartile. + et_summary["selection_rate_by_length_quartile"] = _rate_by_bin( + sub, lambda r: r["len_chars"]) + et_summary["selection_rate_by_offset_quartile"] = _rate_by_bin( + sub, lambda r: r["offset_frac"]) + + # 5. Within-batch position permutation test. + et_summary["within_batch_position"] = _position_permutation_test(sub) + + # 8. Chosen subtype distribution for equation spans. + eq_sel = [r for r in sub if r["span_type"] in EQUATION_TYPES and r["selected"]] + if eq_sel: + et_summary["equation_subtype_choice"] = { + st: dict(Counter(r["chosen_error"] for r in eq_sel if r["span_type"] == st)) + for st in sorted({r["span_type"] for r in eq_sel}) + } + + summary["by_error_type"][et] = et_summary + + # 7. Pool truncation from the 20-per-type extraction cap. + trunc: dict[str, dict[str, int]] = defaultdict(lambda: {"uncapped": 0, "capped": 0, "runs_capped": 0}) + for s in run_stats: + for t, c in s["truncation"].items(): + trunc[t]["uncapped"] += c["uncapped"] + trunc[t]["capped"] += c["capped"] + trunc[t]["runs_capped"] += int(c["uncapped"] > c["capped"]) + summary["pool_truncation_by_span_type"] = { + t: {**v, "fraction_discarded": 1 - v["capped"] / v["uncapped"]} + for t, v in sorted(trunc.items()) + } + + args.output_dir.mkdir(parents=True, exist_ok=True) + rows_path = args.output_dir / "candidate_rows.csv" + with rows_path.open("w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=list(rows[0].keys())) + writer.writeheader() + writer.writerows(rows) + (args.output_dir / "run_stats.json").write_text(json.dumps(run_stats, indent=2)) + summary_path = args.output_dir / "summary.json" + summary_path.write_text(json.dumps(summary, indent=2)) + print(f"\nWrote {rows_path}") + print(f"Wrote {summary_path}") + + # Console digest. + for et, s in summary["by_error_type"].items(): + print(f"\n=== {et} (selection rate {s['overall_selection_rate']:.2f}) ===") + for st, v in s["selection_rate_by_span_type"].items(): + print(f" {st:22s} rate {v['rate']:.2f} ({v['n_selected']}/{v['n_pool']})") + d = s.get("span_type_top2_diff") + if d: + print(f" top-2 type rate diff {d['diff']:+.2f} CI [{d['ci_lo']:+.2f}, {d['ci_hi']:+.2f}]") + lq = s["selection_rate_by_length_quartile"] + print(" by length quartile: " + " ".join(f"{k}={v['rate']:.2f}" for k, v in lq.items())) + oq = s["selection_rate_by_offset_quartile"] + print(" by offset quartile: " + " ".join(f"{k}={v['rate']:.2f}" for k, v in oq.items())) + wb = s["within_batch_position"] + if wb: + print(f" within-batch position shift {wb['observed_mean_position_shift']:+.2f} " + f"(p={wb['p_value']:.4f}, {wb['n_partial_batches']} partial batches)") + + +if __name__ == "__main__": + main()