D19: mask same-passage negatives in InfoNCE
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
3 changed files +18 −3
modified
docs/DECISION_LOG.md
+4 −0
@@ -41,3 +41,7 @@ PREREGISTRATION v1.0 frozen at commit `e74b71266d5c0e8c20f181a50df139f52c771d0f` | ||
| 41 | 41 | ## D18 — 2026-11-09 (pilot diagnostics) |
| 42 | 42 | |
| 43 | 43 | **Finding.** Every M3 variant that included the factual term L_fact (weight 0.5, normalised by the number of protected/numeric tokens) failed to generalise (val CE ≈ 5.4 flat while train CE fell), and so did M1 + L_fact alone (`pilot_m1_fact_tiny`), whereas M3 token-only reached val CE 3.80 ≈ M1 (3.69). The culprit is therefore the normalisation of L_fact, not the dual-latent architecture nor the style/adversarial/contrastive terms. **Decision.** L_fact is now normalised by all target tokens (an extra weight λ_fact on protected tokens, bounded); presets unchanged. Seed variance on the pilot is large (M1 val CE 3.69 / 4.09 / 4.31 over seeds 1–3; M2 seed 2 = 5.03), so pilot-scale comparisons between architectures are not conclusive — the Medium stage is required. |
| 44 | + | |
| 45 | +## D19 — 2026-11-09 (single-term diagnostics) | |
| 46 | + | |
| 47 | +**Finding.** L_contrastive alone (`pilot_m3_cononly_tiny`, val CE stuck ≈ 5.3) also prevented generalisation, L_semantic alone did not (best 4.08). Cause: six rewrites of the same original co-occur in length-bucketed batches and InfoNCE pushed their content vectors apart. **Decision.** Same-passage items are masked out of the InfoNCE negatives (`passage_ids` carried in batches). Presets unchanged; M3 and ablations are re-run as v3 after this fix. | |
modified
src/htr/models/losses.py
+12 −2
@@ -51,10 +51,16 @@ def token_nll(logp: torch.Tensor, tgt_out: torch.Tensor, pad_id: int, smoothing: | ||
| 51 | 51 | return nll.sum(), int(mask.sum().item()) |
| 52 | 52 | |
| 53 | 53 | |
| 54 | −def info_nce(a: torch.Tensor, b: torch.Tensor, temperature: float) -> torch.Tensor: | |
| 54 | +def info_nce(a: torch.Tensor, b: torch.Tensor, temperature: float, same_group: torch.Tensor | None = None) -> torch.Tensor: | |
| 55 | + """Symmetric InfoNCE. ``same_group`` (B, B) bool marks pairs that come from the SAME original passage (several | |
| 56 | + rewrite variants of one passage co-occur in a batch): they are excluded from the negatives instead of being pushed | |
| 57 | + apart (D19 — treating them as negatives corrupted the content encoder in the pilot).""" | |
| 55 | 58 | a = F.normalize(a, dim=-1) |
| 56 | 59 | b = F.normalize(b, dim=-1) |
| 57 | 60 | logits = a @ b.t() / temperature |
| 61 | + if same_group is not None: | |
| 62 | + off_diag_same = same_group & ~torch.eye(a.size(0), dtype=torch.bool, device=a.device) | |
| 63 | + logits = logits.masked_fill(off_diag_same, float("-inf")) | |
| 58 | 64 | labels = torch.arange(a.size(0), device=a.device) |
| 59 | 65 | return 0.5 * (F.cross_entropy(logits, labels) + F.cross_entropy(logits.t(), labels)) |
| 60 | 66 | |
@@ -102,7 +108,11 @@ def compute_loss(model, batch: dict, w: LossWeights) -> tuple[torch.Tensor, dict | ||
| 102 | 108 | enc_t = model.encode(batch["tgt_full"], batch.get("meta")) |
| 103 | 109 | c_t = enc_t["content"].detach() |
| 104 | 110 | if w.con > 0: |
| 105 | − l_con = info_nce(out["content"], c_t, w.temperature) | |
| 111 | + same = None | |
| 112 | + if "passage_ids" in batch: | |
| 113 | + pids = batch["passage_ids"] | |
| 114 | + same = torch.tensor([[p == q for q in pids] for p in pids], dtype=torch.bool, device=out["content"].device) | |
| 115 | + l_con = info_nce(out["content"], c_t, w.temperature, same) | |
| 106 | 116 | logs["con"] = float(l_con.detach()) |
| 107 | 117 | total = total + w.con * l_con |
| 108 | 118 | if w.sem > 0: |
modified
src/htr/training/data.py
+2 −1
@@ -76,7 +76,7 @@ class PairDataset(torch.utils.data.Dataset): | ||
| 76 | 76 | "src": src, "tgt_in": tgt[:-1], "tgt_out": tgt[1:], "fact_mask": fact, "tgt_full": tgt[1:], |
| 77 | 77 | "disc": DISC_IDX.get(m.get("discipline") or "unknown", DISC_IDX["unknown"]), "era": ERA_IDX.get(m.get("era") or "unknown", ERA_IDX["unknown"]), |
| 78 | 78 | "role": ROLE_IDX.get(m.get("section_role") or "other", ROLE_IDX["other"]), "src_id": SRC_IDX.get(m.get("source") or "local", SRC_IDX["local"]), |
| 79 | − "style_label": self.style_map.index(p) if self.style_map else 0, "pair_id": p["pair_id"], | |
| 79 | + "style_label": self.style_map.index(p) if self.style_map else 0, "pair_id": p["pair_id"], "passage_id": p.get("passage_id", p["pair_id"]), | |
| 80 | 80 | } |
| 81 | 81 | |
| 82 | 82 | |
@@ -99,6 +99,7 @@ def collate(items: list[dict], pad_id: int = 0) -> dict: | ||
| 99 | 99 | "meta": {"disc": torch.tensor([x["disc"] for x in items]), "era": torch.tensor([x["era"] for x in items]), "role": torch.tensor([x["role"] for x in items]), "src": torch.tensor([x["src_id"] for x in items])}, |
| 100 | 100 | "style_label": torch.tensor([x["style_label"] for x in items]), |
| 101 | 101 | "pair_ids": [x["pair_id"] for x in items], |
| 102 | + "passage_ids": [x.get("passage_id", x["pair_id"]) for x in items], | |
| 102 | 103 | } |
| 103 | 104 | |
| 104 | 105 | |
| 105 | 106 | |