diff --git a/openspec/changes/ts-as-modality/tasks.md b/openspec/changes/ts-as-modality/tasks.md index cf23f5b..0c46463 100644 --- a/openspec/changes/ts-as-modality/tasks.md +++ b/openspec/changes/ts-as-modality/tasks.md @@ -24,7 +24,7 @@ ## 3. M3 · 阶段①对齐训练 -- [ ] 3.1 损失函数 `train/losses.py`:`lm_loss`(仅回答段计 loss)、`contrastive_loss`(扰动时序→回答 embedding 变化,InfoNCE/扰动一致性),总 loss=`lm_loss+λ*contrastive_loss`(λ 默认 0.1)(验证:loss 标量可反传,对比损失对扰动敏感) +- [x] 3.1 损失函数 `train/losses.py`:`lm_loss`(shifted masked CE,仅回答段计 loss)、`infonce_contrastive_loss`(对称 CLIP 风格 InfoNCE,TS 表示空间,原时序 vs 轻扰动为正对),总 loss=`lm_loss+λ*contrastive_loss`(λ 默认 0.1,由 stage1 组合)(验证:9 单测过——标量有限、全 mask 返 0、shift 正确、grad 可反传、相同对低 loss、对称、L2 不变性)。spec 澄清:InfoNCE 在 TS 嵌入空间(非回答嵌入)做,轻扰动作正对;T3.3 强扰动检验另行评估。 - [ ] 3.2 阶段①训练脚本 `train/stage1.py`+`scripts/train_stage1.sh`:加载 `align.jsonl` 流式,`freeze_llm=True` 仅训 Encoder+Projector(AdamW, lr=1e-4),bs=8/grad_accum=4/ctx=512/BF16/梯度检查点,每 2000 step 存 ckpt + TensorBoard(验证:`--max_steps 100` 冒烟不 OOM、loss 下降、峰值显存 ≤6GB) - [ ] 3.3 TS-token 有效性检验:held-out 原时序 vs 扰动时序回答变化率 + 「必看时序」样本答对率对照纯 LLM(验证:扰动后变化率 >50%,必看时序答对率显著高于纯 LLM) - [ ] 3.4 M3 出口验证:阶段① ckpt 产出;扰动检验通过;峰值显存 ≤6GB(失败回查对比损失权重/数据质量) diff --git a/src/tsmm/train/losses.py b/src/tsmm/train/losses.py new file mode 100644 index 0000000..21bf7d1 --- /dev/null +++ b/src/tsmm/train/losses.py @@ -0,0 +1,70 @@ +"""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) diff --git a/tests/test_losses.py b/tests/test_losses.py new file mode 100644 index 0000000..686f3d8 --- /dev/null +++ b/tests/test_losses.py @@ -0,0 +1,80 @@ +"""Tests for train/losses.py (T3.1).""" +import pytest +import torch + +from tsmm.train.losses import lm_loss, infonce_contrastive_loss + + +class TestLMLoss: + def test_scalar_and_finite(self): + logits = torch.randn(2, 10, 100) + labels = torch.full((2, 10), -100, dtype=torch.long) + labels[:, 5:] = torch.randint(0, 100, (2, 5)) + loss = lm_loss(logits, labels) + assert loss.dim() == 0 + assert torch.isfinite(loss) + + def test_all_masked_returns_zero(self): + logits = torch.randn(2, 10, 100) + labels = torch.full((2, 10), -100, dtype=torch.long) + loss = lm_loss(logits, labels) + assert float(loss) == 0.0 + + def test_shift_correctness(self): + # CE with shift: logits[t] predicts labels[t+1] + torch.manual_seed(0) + V = 5 + logits = torch.zeros(1, 3, V) + # make position 0 strongly predict token 2 + logits[0, 0, 2] = 10.0 + labels = torch.full((1, 3), -100, dtype=torch.long) + labels[0, 1] = 2 # label at position 1 (shifted: predicted from logits[0]) + loss = lm_loss(logits, labels) + assert float(loss) < 0.01 # near-zero when prediction is confident & correct + + def test_gradient_flows(self): + logits = torch.randn(2, 8, 50, requires_grad=True) + labels = torch.full((2, 8), -100, dtype=torch.long) + labels[:, 4:] = torch.randint(0, 50, (2, 4)) + loss = lm_loss(logits, labels) + loss.backward() + assert logits.grad is not None + assert logits.grad.abs().sum() > 0 + + +class TestInfoNCE: + def test_scalar_finite(self): + z1 = torch.randn(4, 16) + z2 = torch.randn(4, 16) + loss = infonce_contrastive_loss(z1, z2) + assert loss.dim() == 0 + assert torch.isfinite(loss) + + def test_identical_pairs_low_loss(self): + torch.manual_seed(0) + z = torch.randn(8, 32) + z = torch.nn.functional.normalize(z, dim=-1) + loss_same = infonce_contrastive_loss(z, z, temperature=0.1) + loss_rand = infonce_contrastive_loss(z, torch.randn(8, 32), temperature=0.1) + assert float(loss_same) < float(loss_rand) + + def test_symmetric(self): + z1 = torch.randn(6, 24) + z2 = torch.randn(6, 24) + a = infonce_contrastive_loss(z1, z2) + b = infonce_contrastive_loss(z2, z1) + assert torch.allclose(a, b, atol=1e-5) + + def test_gradient_flows(self): + z1 = torch.randn(4, 16, requires_grad=True) + z2 = torch.randn(4, 16, requires_grad=True) + loss = infonce_contrastive_loss(z1, z2) + loss.backward() + assert z1.grad is not None and z2.grad is not None + + def test_l2_normalized(self): + # loss should be invariant to per-sample scale (normalize inside) + z1 = torch.randn(4, 16) + loss_a = infonce_contrastive_loss(z1, z1) + loss_b = infonce_contrastive_loss(z1 * 5.0, z1 * 5.0) + assert torch.allclose(loss_a, loss_b, atol=1e-4)