Repository Public Research preview
1"""Training objective L = λ_ce L_token + λ_sem L_semantic + λ_style L_style + λ_con L_contrastive + λ_fact L_factual2(+ the gradient-reversal adversary, which enters through L_adv with the reversed gradient).34Definitions (all differentiable; see docs/architecture/RESEARCH_PLAN.md §Objective):5 L_token label-smoothed negative log-likelihood of target tokens under the (copy-augmented) output distribution.6 L_factual the same NLL restricted to *protected* target tokens (placeholders ⟦C/M/U⟧ and digit-bearing pieces),7 i.e. an up-weighting of factual spans; weight 0 → plain token loss.8 L_style cross-entropy of the style branch E_s against the metadata-derived style class (training-only labels).9 L_adv cross-entropy of the adversary that predicts the style class from the *content* vector through a10 gradient-reversal layer; minimising it w.r.t. the adversary while the reversed gradient pushes the11 content encoder to be style-uninformative.12 L_contrastive InfoNCE between the content vector of the rewrite and the content vector of the original (encoded13 by the same encoder with stop-gradient), negatives = other items in the batch.14 L_semantic 1 − cosine(content(rewrite), content(original)) (stop-gradient on the original side); cheaper15 surrogate of an external embedding similarity.16 L_orth squared Frobenius norm of the cross-covariance between content and style vectors (batch estimate).17Weights are hyper-parameters tuned on validation only; the token-loss-only configuration is the M1/M2 control.18"""19from __future__ import annotations2021from dataclasses import dataclass2223import torch24import torch.nn.functional as F252627@dataclass28class LossWeights:29 ce: float = 1.030 fact: float = 0.031 style: float = 0.032 adv: float = 0.033 con: float = 0.034 sem: float = 0.035 orth: float = 0.036 label_smoothing: float = 0.137 temperature: float = 0.07383940def token_nll(logp: torch.Tensor, tgt_out: torch.Tensor, pad_id: int, smoothing: float, weights: torch.Tensor | None = None) -> tuple[torch.Tensor, int]:41 """logp (B, L, V) log-probs; tgt_out (B, L). Returns (sum loss, n_tokens)."""42 V = logp.size(-1)43 mask = tgt_out.ne(pad_id)44 nll = -logp.gather(-1, tgt_out.unsqueeze(-1)).squeeze(-1)45 if smoothing > 0:46 smooth = -logp.mean(-1)47 nll = (1 - smoothing) * nll + smoothing * smooth48 if weights is not None:49 nll = nll * weights50 nll = nll.masked_fill(~mask, 0.0)51 return nll.sum(), int(mask.sum().item())525354def 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 (several56 rewrite variants of one passage co-occur in a batch): they are excluded from the negatives instead of being pushed57 apart (D19 — treating them as negatives corrupted the content encoder in the pilot)."""58 a = F.normalize(a, dim=-1)59 b = F.normalize(b, dim=-1)60 logits = a @ b.t() / temperature61 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"))64 labels = torch.arange(a.size(0), device=a.device)65 return 0.5 * (F.cross_entropy(logits, labels) + F.cross_entropy(logits.t(), labels))666768def orthogonality(c: torch.Tensor, s: torch.Tensor) -> torch.Tensor:69 c = c - c.mean(0, keepdim=True)70 s = s - s.mean(0, keepdim=True)71 cov = c.t() @ s / max(1, c.size(0) - 1)72 return (cov**2).sum() / (c.size(1) * s.size(1))737475def compute_loss(model, batch: dict, w: LossWeights) -> tuple[torch.Tensor, dict]:76 """batch: src, tgt_in, tgt_out, fact_mask (B,L) float, meta dict, style_label (B) long, tgt_full (B,Lt) for encoding the original."""77 out = model(batch["src"], batch["tgt_in"], batch.get("meta"))78 pad = model.cfg.pad_id79 logs: dict[str, float] = {}80 total = torch.zeros((), device=batch["src"].device)81 nll_sum, n_tok = token_nll(out["logp"], batch["tgt_out"], pad, w.label_smoothing)82 l_ce = nll_sum / max(1, n_tok)83 logs["ce"] = float(l_ce.detach())84 total = total + w.ce * l_ce85 if w.fact > 0 and "fact_mask" in batch:86 fw = batch["fact_mask"]87 f_sum, _ = token_nll(out["logp"], batch["tgt_out"], pad, 0.0, weights=fw)88 # Normalise by ALL target tokens (same denominator as L_token): L_fact is then an extra weight of λ_fact on the89 # protected/numeric tokens. Normalising by the number of protected tokens (pilot, D18) let a handful of tokens90 # dominate every update and destroyed generalisation (val CE stuck near 5.4 while train CE fell).91 l_fact = f_sum / max(1, n_tok)92 logs["fact"] = float(l_fact.detach())93 logs["fact_per_token"] = float((f_sum / fw.sum().clamp(min=1.0)).detach())94 total = total + w.fact * l_fact95 if model.cfg.arch == "m3":96 if w.style > 0 and "style_label" in batch:97 l_style = F.cross_entropy(out["style_logits"], batch["style_label"])98 logs["style"] = float(l_style.detach())99 logs["style_acc"] = float((out["style_logits"].argmax(-1) == batch["style_label"]).float().mean())100 total = total + w.style * l_style101 if w.adv > 0 and "style_label" in batch:102 l_adv = F.cross_entropy(out["adv_logits"], batch["style_label"])103 logs["adv"] = float(l_adv.detach())104 logs["adv_acc"] = float((out["adv_logits"].argmax(-1) == batch["style_label"]).float().mean())105 total = total + w.adv * l_adv106 if (w.con > 0 or w.sem > 0) and "tgt_full" in batch:107 with torch.no_grad():108 enc_t = model.encode(batch["tgt_full"], batch.get("meta"))109 c_t = enc_t["content"].detach()110 if w.con > 0:111 same = None112 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)116 logs["con"] = float(l_con.detach())117 total = total + w.con * l_con118 if w.sem > 0:119 l_sem = 1 - F.cosine_similarity(out["content"], c_t, dim=-1).mean()120 logs["sem"] = float(l_sem.detach())121 total = total + w.sem * l_sem122 if w.orth > 0:123 l_orth = orthogonality(out["content"], out["style_vec"])124 logs["orth"] = float(l_orth.detach())125 total = total + w.orth * l_orth126 logs["total"] = float(total.detach())127 logs["n_tok"] = n_tok128 return total, logs129