File size: 1,493 Bytes
5952424 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | """Multi-head loss: char-diffusion loss dominates; boundary/diacritics/cap/punctuation are
down-weighted auxiliary heads, jointly trained."""
from __future__ import annotations
import torch
import torch.nn.functional as F
def _anchored_ce(logits, labels, weight, logs, name):
"""Cross-entropy on supervised positions; if none, return a ZERO term that still touches
every head param so DDP never sees an 'unused parameter'."""
mm = labels != -100
if mm.any():
l = F.cross_entropy(logits[mm], labels[mm])
logs[name] = l.item()
return weight * l
return logits.sum() * 0.0 # zero grad, but head is "used"
def compute_loss(out, batch, lam=0.1):
logs = {}
# char diffusion loss with MDLM 1/t reweighting
cl = out["char"]
lab = batch["labels"]
m = lab != -100
if m.any():
ce = F.cross_entropy(cl[m], lab[m], reduction="none")
w = batch["loss_w"][m]
char_loss = (ce * w).sum() / w.sum()
else:
char_loss = cl.sum() * 0.0
logs["char"] = char_loss.item()
total = char_loss
total = total + _anchored_ce(out["boundary"], batch["bnd_lab"], lam, logs, "bnd")
total = total + _anchored_ce(out["dia"], batch["dia_lab"], lam, logs, "dia")
total = total + _anchored_ce(out["cap"], batch["cap_lab"], lam, logs, "cap")
total = total + _anchored_ce(out["punct"], batch["punct_lab"], lam, logs, "punct")
logs["total"] = total.item()
return total, logs
|