Buckets:
| """E3 losses: task loss + (on items with a cached teacher) L1 between LayerNorm-normalised student and teacher | |
| option-slot hidden states at blocks 14 and 28 (final pass, backbone state before the head) + KL(p_t || p_s).""" | |
| import torch | |
| import torch.nn.functional as F | |
| from model import collate, item_logits | |
| from shared_yesno_supervised import supervised_terms | |
| def attach_teacher(items, path): | |
| cache = torch.load(path, weights_only=False) | |
| for x in items: | |
| t = cache.get("%s|%s" % (x["case_id"], x["view_id"])) | |
| if t is not None: | |
| assert t["h14"].shape[0] == len(x["pos"]), "teacher/student option slots differ" | |
| x["teacher"] = t | |
| return cache | |
| def distill_losses(model, items, K, device, pad_id, a): | |
| layers = tuple(int(l) for l in a.distill_layers.split(",")) | |
| batch = collate(items, device, pad_id) | |
| outs, caps = model(batch, K=K, capture=layers) | |
| logits = item_logits(outs[-1], batch) | |
| terms = [] | |
| l1_sum = kl_sum = 0.0 | |
| extra = torch.zeros((), device=device) | |
| for j, (lg, x) in enumerate(zip(logits, items)): | |
| t = torch.tensor([x["target"]], dtype=torch.float32, device=device) | |
| terms.append(supervised_terms(lg, t, x["primitive"], a.rps_weight)) | |
| tch = x.get("teacher") | |
| if tch is None: | |
| continue | |
| n = len(x["pos"]) | |
| l1 = 0.0 | |
| for l in layers: | |
| hs = F.layer_norm(caps[l][j, :n].float(), (caps[l].shape[-1],)) | |
| ht = tch["h%d" % l].to(device, torch.float32) | |
| l1 = l1 + (hs - ht).abs().mean() | |
| pt = tch["p"].to(device) | |
| logp_s = F.log_softmax(lg[0].float(), -1) | |
| kl = (pt * (torch.log(pt.clamp_min(1e-12)) - logp_s)).sum() | |
| extra = extra + a.distill_l1 * l1 + a.distill_kl * kl | |
| l1_sum += float(l1.detach()) if torch.is_tensor(l1) else float(l1) | |
| kl_sum += float(kl.detach()) | |
| return terms, dict(loss=extra, l1=l1_sum, dkl=kl_sum) | |
Xet Storage Details
- Size:
- 1.95 kB
- Xet hash:
- 2e44ef77ca195f45d67098d7959a297db98bb7e7df2972ec644066e1e382dbf2
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.