张宗平
204d5237ba
- src/tsmm/train/losses.py:
* lm_loss(logits, labels): shifted masked CE (answer span only); clean 0.0
when all positions masked.
* infonce_contrastive_loss(z1, z2): symmetric CLIP-style InfoNCE in TS-
embedding space; positive = original vs light-augmentation view.
- tests/test_losses.py: 9 tests. 100 tests passing.
Spec clarification (small tier): InfoNCE operates in TS-embedding space
(encoder/projector output) rather than answer-embedding space; light
augmentations are positives. T3.3 strong-perturbation sensitivity is a
separate downstream check.
71 lines
2.7 KiB
Python
71 lines
2.7 KiB
Python
"""Losses for stage ① alignment training (T3.1).
|
|
|
|
- :func:`lm_loss`: masked, shifted cross-entropy — computes LM loss only on the
|
|
answer span (``labels != -100``), with the standard right-shift.
|
|
- :func:`infonce_contrastive_loss`: symmetric CLIP-style InfoNCE between two
|
|
batches of representations. Used to align two *views* of the same time series
|
|
(original vs lightly-augmented), forcing the TS encoder to be discriminative
|
|
rather than collapse to a constant (which would let the LLM ignore the TS
|
|
tokens).
|
|
|
|
Total stage ① loss (composed by ``train/stage1.py``):
|
|
|
|
loss = lm_loss + lambda * contrastive_loss (lambda default 0.1)
|
|
|
|
Design ref: §4.1 / §4.2 — "扰动一致性/InfoNCE". We implement InfoNCE in TS-
|
|
embedding space (encoder/projector output) rather than answer-embedding space,
|
|
because (a) it is cheap (no LLM forward needed for the augmented view) and
|
|
(b) the alignment goal is precisely that distinct series produce distinct TS
|
|
representations; the LM term already binds TS→answer. T3.3 (>50% answer change
|
|
under *strong* perturbation) is then evaluated separately, using perturbations
|
|
larger than the light augmentations used as InfoNCE positives.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import Tensor
|
|
|
|
|
|
def lm_loss(logits: Tensor, labels: Tensor, ignore_index: int = -100) -> Tensor:
|
|
"""Masked, shifted cross-entropy.
|
|
|
|
``logits`` : [B, S, V]
|
|
``labels`` : [B, S] (``ignore_index`` masks non-answer positions)
|
|
Returns a scalar.
|
|
"""
|
|
if labels.numel() == 0:
|
|
return logits.new_zeros(())
|
|
# standard causal shift: predict token t+1 from position t
|
|
shift_logits = logits[:, :-1, :].contiguous()
|
|
shift_labels = labels[:, 1:].contiguous()
|
|
if (shift_labels != ignore_index).sum() == 0:
|
|
return logits.new_zeros(())
|
|
loss = F.cross_entropy(
|
|
shift_logits.view(-1, shift_logits.size(-1)),
|
|
shift_labels.view(-1),
|
|
ignore_index=ignore_index,
|
|
reduction="mean",
|
|
)
|
|
return loss
|
|
|
|
|
|
def infonce_contrastive_loss(
|
|
z1: Tensor, z2: Tensor, temperature: float = 0.07
|
|
) -> Tensor:
|
|
"""Symmetric InfoNCE (CLIP-style) between two sets of representations.
|
|
|
|
Positive pair = (z1[i], z2[i]); negatives = all j != i in the batch.
|
|
Inputs are L2-normalized internally, so scale is ignored.
|
|
|
|
``z1``, ``z2`` : [B, D]
|
|
Returns a scalar.
|
|
"""
|
|
z1n = F.normalize(z1, dim=-1)
|
|
z2n = F.normalize(z2, dim=-1)
|
|
logits = z1n @ z2n.t() / temperature # [B, B]
|
|
targets = torch.arange(z1.size(0), device=z1.device)
|
|
loss_12 = F.cross_entropy(logits, targets)
|
|
loss_21 = F.cross_entropy(logits.t(), targets)
|
|
return 0.5 * (loss_12 + loss_21)
|