Theme-overlap analysis: verify train/eval split is not theme-contaminated

[sebasmos]

Why this matters for the paper

Deduplication by prompt_id does not prevent shared rubric themes (emergency_referrals, hedging, context_seeking) from appearing in both train and eval. A reviewer will ask whether HealthBench score gains reflect rubric pattern memorization rather than generalization. Without checking tag distribution across splits, the evaluation cannot distinguish learning from leakage.

Scope

• Extend scripts/check_dataset_overlap.py to report example_tags distribution across train (data/sft/train.jsonl) vs eval (data/raw/hard_200_sample_ids.json) splits, not just prompt_id overlap.
• Report tag frequency counts and flag any tag that appears >2x more in train than eval (or vice versa).
• If the split is unbalanced, stratify the 200-example holdout by theme and report per-theme scores in the eval pipeline.

Done when

check_dataset_overlap.py prints a tag-frequency table and a list of imbalanced tags. If imbalance is found, per-theme eval scores are reported in the results JSON.

Effort

Small — extend existing script, ~2h.

Good for

sahilk873 or whoever runs the next grading pass.

[sebasmos]

check_dataset_overlap.py now includes tag-overlap checking.

The --tag-overlap flag compares tag/theme frequency between the training pool and each eval draw. Run:

python scripts/check_dataset_overlap.py \
    --train-jsonl data/sft/train.jsonl \
    --seeds 5 \
    --tag-overlap

Any tag with >2x frequency skew between train and eval is flagged. If a theme like emergency_referrals or hedging is heavily over-represented in training relative to the eval draw, scores on that theme are not a fair measure of generalization.

The core ask of this issue (extending the script) is done. What remains:

• Run the script once the raw data files are downloaded and data/sft/train.jsonl exists.
• If imbalances are found, decide whether to stratify the eval draw or report per-theme scores separately.

[sebasmos]

Cross-reference: the per-seed draw protocol is defined in #3.

This issue focuses on tag/theme overlap analysis. The seed draws it should check are the 5 independent 200-sample subsets from the 1K Hard set described in #3. Run scripts/check_dataset_overlap.py --seeds 5 --tag-overlap --train-jsonl data/sft/train.jsonl to check both leakage and tag balance across all five draws simultaneously.

[sebasmos]

Team allocation: tag-overlap analysis is partly automated, partly manual. The manual review of which tags should count as 'safety-themes' for the per-theme ablation (cross-references #18) is annotation work — Peter/Sahil should assign Max (maximinl), Ashley, Martha, and Sohyeon. Add their handles. Clinical judgment on whether a tag is medically meaningful would benefit from input by Zineb, Ash Doulla, Hillary, or Felipe.

[maximinl]

Bumping this; after digging through the pipeline today, the issue is bigger than the tag-distribution check this ticket was originally scoped to. The structural problem:

Train pool ⊃ bootstrap eval pool, by ~80% per seed.

• tpu/launch_stage1.sh:143 and slurm/generate_traces.sh:41 only exclude the fixed 200-prompt holdout (hard_200_sample_ids.json), not all of Hard. The other ~800 Hard prompts stay in raw_traces.jsonl.
• run_multi_seed.sh reuses one raw_traces.jsonl across all 5 seeds, and filter_traces.py has no --exclude-ids parameter, so those ~800 Hard prompts flow into every seed's data/sft/seed_*/train.jsonl.
• Each per-seed bootstrap draws 200 random IDs from 1000 Hard. Expected overlap with train = 200 × 800/1000 = ~160 of 200 prompts (~80%) per seed.

RESULTS.md:29 claims "for each seed, the 200-sample draw must be excluded from training data generation" but no script implements this — it would require regenerating traces 5 times. The headline multi-seed numbers are in-distribution, not held-out.

Two paths, and we need to pick one before the next grading pass:

1. Re-train clean. Regenerate raw_traces.jsonl with --exclude-ids data/raw/healthbench_hard.jsonl (drops all 1K Hard) → train on Full minus Hard (~4000 prompts) → the per-seed bootstrap from Hard is then honest. Costs ~40h of H100 per the comment in run_multi_seed.sh:4. Cleanest for the paper.
2. Keep the current run, justify in writing. Report current numbers as in-distribution; report a separate generalization number from the existing slurm/generalization.sh artifact (which already excludes all of Hard correctly). Methods section needs a paragraph naming the contamination and saying what each number measures. No new compute.

@PeterLi-jpg @sahilk873 — which one?

In parallel I'll open a small PR wiring scripts/check_dataset_overlap.py into run_multi_seed.sh as a preflight gate that exits non-zero on any leakage, so this can't ship silently again. That's independent of which path we pick.