consistent train/eval template split

Author: David Johnston

examples/semantic/asymmetric.py

@@ -45,6 +45,9 @@ class AsymmetricConfig:
     seed: int = 42
     # HuggingFace dataset repo. If set, skips local generation and downloads from HF.
     hf_dataset: str | None = None
+    # Template split for train/test segregation (only used for local generation)
+    # Train uses templates < train_template_cutoff, eval majority uses templates >= cutoff
+    train_template_cutoff: int = 5
 
 
 def create_asymmetric_dataset(
@@ -57,6 +60,11 @@ def create_asymmetric_dataset(
     - Exclusive facts: only appear in dominant style (for testing semantic matching)
     - Shared facts: appear in both styles (for style ratio control)
 
+    For train/test segregation:
+    - Training uses templates < train_template_cutoff (default: 0-4)
+    - Eval majority style control uses templates >= cutoff (default: 5+)
+    This ensures no exact text overlap between train and eval majority control.
+
     Args:
         config: Experiment configuration.
         output_dir: Directory to save datasets.
@@ -103,29 +111,41 @@ def create_asymmetric_dataset(
     dominant_ds = style_datasets[config.dominant_style]
     minority_ds = style_datasets[config.minority_style]
 
-    # Get unique facts (by fact text)
-    all_facts = list(set(original["fact"]))
-    n_facts = len(all_facts)
+    # Get unique (identifier, field) pairs - these represent underlying semantic facts
+    # Each pair has multiple templates (different surface forms of the same fact)
+    semantic_facts = list(
+        {(row["identifier"], row["field"]) for row in original}
+    )
+    n_semantic_facts = len(semantic_facts)
 
-    # Split into exclusive (dominant-only) and shared
+    # Split into exclusive (dominant-only) and shared by semantic fact
     rng = np.random.default_rng(config.seed)
-    rng.shuffle(all_facts)
-
-    n_exclusive = int(n_facts * config.exclusive_ratio)
-    exclusive_facts = set(all_facts[:n_exclusive])
-    shared_facts = set(all_facts[n_exclusive:])
-
-    print(f"Total unique facts: {n_facts}")
-    print(f"Exclusive to {config.dominant_style}: {len(exclusive_facts)}")
-    print(f"Shared between styles: {len(shared_facts)}")
-
-    # Build training set
-    # 1. All dominant style facts
-    train_dominant = dominant_ds
+    rng.shuffle(semantic_facts)
+
+    n_exclusive = int(n_semantic_facts * config.exclusive_ratio)
+    exclusive_semantic_facts = set(semantic_facts[:n_exclusive])
+    shared_semantic_facts = set(semantic_facts[n_exclusive:])
+
+    print(f"Total unique semantic facts (identifier, field pairs): {n_semantic_facts}")
+    print(f"Exclusive to {config.dominant_style}: {len(exclusive_semantic_facts)}")
+    print(f"Shared between styles: {len(shared_semantic_facts)}")
+    print(f"Template cutoff for train/eval split: {config.train_template_cutoff}")
+
+    # Build training set with template filtering
+    # 1. Dominant style: only templates < cutoff (to reserve rest for eval majority control)
+    train_dominant_indices = [
+        i
+        for i, row in enumerate(dominant_ds)
+        if row["template"] < config.train_template_cutoff
+    ]
+    train_dominant = dominant_ds.select(train_dominant_indices)
 
-    # 2. Minority style only for shared facts
+    # 2. Minority style only for shared facts (any template since minority eval is different)
     minority_shared_indices = [
-        i for i, row in enumerate(minority_ds) if row["fact"] in shared_facts
+        i
+        for i, row in enumerate(minority_ds)
+        if (row["identifier"], row["field"]) in shared_semantic_facts
+        and row["template"] < config.train_template_cutoff
     ]
     train_minority = minority_ds.select(minority_shared_indices)
 
@@ -148,10 +168,14 @@ def create_asymmetric_dataset(
     print(f"  Dominant ratio: {len(train_dominant) / len(train_ds):.2%}")
 
     # Build eval set: query exclusive facts in minority style
+    # Use templates >= cutoff to ensure no overlap with train
     # These facts don't exist in minority style in training, so the model
     # must use semantic matching (not style matching) to find them
     eval_minority_indices = [
-        i for i, row in enumerate(minority_ds) if row["fact"] in exclusive_facts
+        i
+        for i, row in enumerate(minority_ds)
+        if (row["identifier"], row["field"]) in exclusive_semantic_facts
+        and row["template"] >= config.train_template_cutoff
     ]
     eval_ds = minority_ds.select(eval_minority_indices)
     eval_ds = eval_ds.add_column("style", [config.minority_style] * len(eval_ds))
@@ -164,6 +188,7 @@ def create_asymmetric_dataset(
     print("\nEval set:")
     print(f"  Queries in {config.minority_style} style: {len(eval_ds)}")
     print(f"  Ground truth only in {config.dominant_style} style")
+    print(f"  Using templates >= {config.train_template_cutoff} (no overlap with train)")
 
     # Save datasets
     train_ds.save_to_disk(str(train_path))
@@ -1048,40 +1073,81 @@ def compute_asymmetric_metrics_with_pca(
 def create_majority_style_eval(
     config: AsymmetricConfig,
     base_path: Path | str,
-) -> Path:
+    force_regenerate: bool = False,
+) -> tuple[Path, bool]:
     """Create eval set using majority style (control for style mismatch).
 
     Instead of using minority style queries, uses dominant style queries
     for the exclusive facts. This shows baseline performance without style mismatch.
 
+    IMPORTANT: Uses templates >= train_template_cutoff to ensure NO overlap with
+    training data. This provides a proper train/test split where eval majority
+    style items test semantic matching (same fact, different surface form) rather
+    than exact text matching.
+
     Args:
         config: Experiment configuration.
         base_path: Base path for experiment outputs.
+        force_regenerate: If True, regenerate even if cached version exists.
 
     Returns:
-        Path to the majority style eval dataset.
+        Tuple of (path to the majority style eval dataset, has_leakage flag).
+        has_leakage is True if there's train/test overlap (e.g., from HF data).
     """
     base_path = Path(base_path)
     data_path = base_path / "data"
     majority_eval_path = data_path / "eval_majority_style.hf"
 
-    if majority_eval_path.exists():
+    # Check for existing cached version
+    if majority_eval_path.exists() and not force_regenerate:
         print(f"Loading cached majority style eval from {majority_eval_path}")
-        return majority_eval_path
+
+        # Check for train/test leakage by comparing reworded texts
+        train_ds = load_from_disk(str(data_path / "train.hf"))
+        majority_eval_ds = load_from_disk(str(majority_eval_path))
+        if isinstance(train_ds, DatasetDict):
+            train_ds = train_ds["train"]
+        if isinstance(majority_eval_ds, DatasetDict):
+            majority_eval_ds = majority_eval_ds["train"]
+
+        train_reworded = set(train_ds["reworded"])
+        eval_reworded = set(majority_eval_ds["reworded"])
+        overlap = train_reworded & eval_reworded
+        has_leakage = len(overlap) > 0
+
+        if has_leakage:
+            print(
+                f"  WARNING: {len(overlap)}/{len(eval_reworded)} eval items have "
+                "exact text match in train (train/test leakage)"
+            )
+            print("  Use force_regenerate=True with local data to fix")
+
+        return majority_eval_path, has_leakage
 
     print("Creating majority style eval set (control)...")
 
-    # Load the minority style eval to get the facts
+    # Check if local styled datasets exist for proper template segregation
+    local_styled_path = Path(
+        f"data/facts_dataset_{config.dominant_style}-Qwen3-8B-Base.hf"
+    )
+    if not local_styled_path.exists():
+        print(
+            f"  WARNING: Local styled dataset not found at {local_styled_path}"
+        )
+        print("  Cannot create properly segregated majority eval")
+        print("  Using HF eval_majority_style (may have train/test leakage)")
+        return majority_eval_path, True  # Return existing HF version with leakage flag
+
+    # Load the minority style eval to get the semantic facts (identifier, field pairs)
     eval_ds = load_from_disk(str(data_path / "eval.hf"))
     if isinstance(eval_ds, DatasetDict):
         eval_ds = eval_ds["train"]
 
-    eval_facts = set(eval_ds["fact"])
+    # Get semantic facts from eval (identifier, field pairs)
+    eval_semantic_facts = {(row["identifier"], row["field"]) for row in eval_ds}
 
     # Load dominant style dataset
-    dominant_ds = load_from_disk(
-        f"data/facts_dataset_{config.dominant_style}-Qwen3-8B-Base.hf"
-    )
+    dominant_ds = load_from_disk(str(local_styled_path))
     if isinstance(dominant_ds, DatasetDict):
         dominant_ds = dominant_ds["train"]
 
@@ -1096,24 +1162,33 @@ def create_majority_style_eval(
             restored_col = [fact_to_meta[row["fact"]][col] for row in dominant_ds]
             dominant_ds = dominant_ds.add_column(col, restored_col)
 
-    # Select only the exclusive facts (same facts as in minority eval)
+    # Select dominant style versions of eval semantic facts
+    # Use templates >= cutoff to ensure NO overlap with training data
     dominant_eval_indices = [
-        i for i, row in enumerate(dominant_ds) if row["fact"] in eval_facts
+        i
+        for i, row in enumerate(dominant_ds)
+        if (row["identifier"], row["field"]) in eval_semantic_facts
+        and row["template"] >= config.train_template_cutoff
     ]
     majority_eval_ds = dominant_ds.select(dominant_eval_indices)
 
-    # Add style columns
-    majority_eval_ds = majority_eval_ds.add_column(
-        "style", [config.dominant_style] * len(majority_eval_ds)
-    )
-    majority_eval_ds = majority_eval_ds.add_column(
-        "expected_match_style", [config.dominant_style] * len(majority_eval_ds)
-    )
+    print(f"  Using templates >= {config.train_template_cutoff} (no overlap with train)")
+    print(f"  Found {len(majority_eval_ds)} majority style eval samples")
+
+    # Add style columns if not present
+    if "style" not in majority_eval_ds.column_names:
+        majority_eval_ds = majority_eval_ds.add_column(
+            "style", [config.dominant_style] * len(majority_eval_ds)
+        )
+    if "expected_match_style" not in majority_eval_ds.column_names:
+        majority_eval_ds = majority_eval_ds.add_column(
+            "expected_match_style", [config.dominant_style] * len(majority_eval_ds)
+        )
 
     majority_eval_ds.save_to_disk(str(majority_eval_path))
     print(f"Saved majority style eval to {majority_eval_path}")
 
-    return majority_eval_path
+    return majority_eval_path, False  # No leakage with proper segregation
 
 
 def score_majority_style_eval(
@@ -1146,7 +1221,9 @@ def score_majority_style_eval(
     data_path = base_path / "data"
 
     # Create majority style eval if needed
-    create_majority_style_eval(config, base_path)
+    _, has_leakage = create_majority_style_eval(config, base_path)
+    if has_leakage:
+        print("  Note: Majority control may show inflated accuracy due to train/test leakage")
 
     # Determine output path
     if preconditioner_name:
@@ -1291,7 +1368,7 @@ def compute_majority_style_metrics(
     data_path = base_path / "data"
 
     # Create majority style eval if needed
-    create_majority_style_eval(config, base_path)
+    _, _ = create_majority_style_eval(config, base_path)
 
     # Load datasets
     train_ds = load_from_disk(str(data_path / "train.hf"))
@@ -1445,7 +1522,7 @@ def score_summed_eval(
         eval_minority_ds = eval_minority_ds["train"]
 
     # Create majority style eval if needed
-    create_majority_style_eval(config, base_path)
+    _, _ = create_majority_style_eval(config, base_path)
     eval_majority_ds = load_from_disk(str(data_path / "eval_majority_style.hf"))
     if isinstance(eval_majority_ds, DatasetDict):
         eval_majority_ds = eval_majority_ds["train"]
@@ -1455,17 +1532,21 @@ def score_summed_eval(
         f"Scoring {n_eval} summed eval queries (minority + majority) against {n_train} train samples"
     )
 
-    # Build fact-to-index mapping for alignment
-    minority_facts = eval_minority_ds["fact"]
-    majority_facts = eval_majority_ds["fact"]
-    majority_fact_to_idx = {f: i for i, f in enumerate(majority_facts)}
+    # Build semantic fact mapping for alignment (identifier, field pairs)
+    # This works even when templates differ between minority and majority eval
+    minority_semantic_facts = [
+        (row["identifier"], row["field"]) for row in eval_minority_ds
+    ]
+    majority_semantic_to_idx = {
+        (row["identifier"], row["field"]): i for i, row in enumerate(eval_majority_ds)
+    }
 
-    # Verify alignment
-    assert len(minority_facts) == len(
-        majority_facts
+    # Verify alignment by semantic fact
+    assert len(eval_minority_ds) == len(
+        eval_majority_ds
     ), "Eval datasets must have same size"
-    for f in minority_facts:
-        assert f in majority_fact_to_idx, f"Fact {f} not found in majority eval"
+    for sf in minority_semantic_facts:
+        assert sf in majority_semantic_to_idx, f"Semantic fact {sf} not found in majority eval"
 
     # Load train gradients
     print("Loading train gradients...")
@@ -1575,14 +1656,16 @@ def score_summed_eval(
     majority_grads = load_gradients(eval_majority_grads_path, structured=True)
 
     # Sum gradients: for each eval fact, sum minority + majority style gradients
-    # Need to align by fact since ordering might differ
+    # Align by semantic fact (identifier, field) since templates may differ
     summed_grad_list = []
     for name in tqdm(module_names, desc="Summing eval grads"):
         g_minority = torch.from_numpy(_load_gradients_as_float(minority_grads, name))
         g_majority = torch.from_numpy(_load_gradients_as_float(majority_grads, name))
 
-        # Align majority grads to minority fact order
-        aligned_majority_indices = [majority_fact_to_idx[f] for f in minority_facts]
+        # Align majority grads to minority semantic fact order
+        aligned_majority_indices = [
+            majority_semantic_to_idx[sf] for sf in minority_semantic_facts
+        ]
         g_majority_aligned = g_majority[aligned_majority_indices]
 
         # Sum the gradients
@@ -2019,6 +2102,21 @@ def run_asymmetric_experiment(
         print("\n" + "-" * 60)
         print("SEMANTIC-ONLY EVAL (gradients only from answer tokens)")
         print("-" * 60)
+
+        # Standard influence function approach: semantic mask + H_train preconditioner
+        # This is the "correct" way to compute influence functions
+        print("\n--- Strategy: semantic_index (standard IF with H_train) ---")
+        metrics = compute_asymmetric_metrics(
+            config,
+            base_path,
+            "index",  # H_train - the standard IF preconditioner
+            damping_factor=damping_factor,
+            eval_prompt_column="question",
+            eval_completion_column="answer",
+        )
+        print_metrics(metrics, "semantic_index")
+        all_metrics["semantic_index"] = metrics
+
         print("\n--- Strategy: semantic_no_precond ---")
         metrics = compute_asymmetric_metrics(
             config,
@@ -2193,8 +2291,14 @@ def score_with_inner_product(
         if isinstance(eval_majority_ds, DatasetDict):
             eval_majority_ds = eval_majority_ds["train"]
 
-        minority_facts = eval_minority_ds["fact"]
-        majority_fact_to_idx = {f: i for i, f in enumerate(eval_majority_ds["fact"])}
+        # Use semantic fact alignment (identifier, field) since templates may differ
+        minority_semantic_facts = [
+            (row["identifier"], row["field"]) for row in eval_minority_ds
+        ]
+        majority_semantic_to_idx = {
+            (row["identifier"], row["field"]): i
+            for i, row in enumerate(eval_majority_ds)
+        }
 
         summed_grad_list = []
         for name in tqdm(module_names, desc="Summing eval grads"):
@@ -2205,7 +2309,9 @@ def score_with_inner_product(
                 _load_gradients_as_float(majority_grads, name)
             )
 
-            aligned_majority_indices = [majority_fact_to_idx[f] for f in minority_facts]
+            aligned_majority_indices = [
+                majority_semantic_to_idx[sf] for sf in minority_semantic_facts
+            ]
             g_majority_aligned = g_majority[aligned_majority_indices]
 
             g_summed = g_minority + g_majority_aligned