feat(m3): losses — masked LM CE + symmetric InfoNCE (T3.1)

- 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.

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（失败回查对比损失权重/数据质量）

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)

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)