[feat] support session_name in listwise_rank_loss and jrc_loss

docs/source/models/loss.md

@@ -103,6 +103,8 @@ model_config {
 适用二分类任务的损失函数，其对应的任务num_class必须是2。该损失函数除了关注样本目标自身分类的正确性，还会关注在同一个batch的同一个session中，所有正样本的概率要尽可能的大于所有负样本的概率。
 https://arxiv.org/abs/2208.06164
 
+session的来源与[listwise_rank_loss](#listwise_rank_loss)相同：设置`session_name`时，同一个batch内该字段取值相同的样本构成一个session（不要求相邻）；不设置时，DlrmHSTU系列模型以每个请求的候选序列为一个session，其他排序模型必须设置`session_name`
+
 配置如下
 
 ```
@@ -115,7 +117,7 @@ model_config {
 }
 ```
 
-对于该损失函数，要求同一个session_id的样本尽量在一个batch中进行训练，在一个session中尽量要求样本保持有序。
+对于该损失函数，要求同一个session_id的样本落在同一个batch中。
 
 我们使用sql如下方式构造样本,该数据集的session_name是user_id
 
@@ -131,14 +133,21 @@ SORT BY user_id asc,time_stamp asc
 
 ## listwise_rank_loss
 
-请求粒度的listwise排序损失（InfoNCE），同一个请求内的其他候选互为负样本。该损失函数要求模型发布每个请求的候选数以及该任务的`logits_<task_name>`预测，目前只有DlrmHSTU系列满足，详见[dlrm_hstu_onerank](dlrm_hstu_onerank.md)。 一个请求需要同时含有至少一个正样本和一个负样本才会计入该项损失
+listwise排序损失（InfoNCE），同一个list内的其他样本互为负样本。list的来源有两种：
+
+- DlrmHSTU系列模型：模型会发布每个请求的候选数，一个请求的候选序列即为一个list，无需额外配置（设置`session_name`会报错），详见[dlrm_hstu_onerank](dlrm_hstu_onerank.md)
+- 其他排序模型（多塔、DBMTL等）：必须通过`session_name`指定分组字段（如request_id或user_id），同一个batch内该字段取值相同的样本构成一个list。该字段需为模型的稀疏特征
 
 配置如下
 
 ```
 model_config {
+    losses {
+        binary_cross_entropy {}
+    }
     losses {
         listwise_rank_loss {
+            session_name: "user_id"
             temperature_init: 0.07
             learnable_temperature: true
         }
@@ -149,10 +158,14 @@ model_config {
 
 参数说明：
 
+1. session_name: list分组的字段名，仅用于非DlrmHSTU系列模型（DlrmHSTU系列模型按请求分组，不可设置）
 1. temperature_init: softmax温度初始值，logits会乘以 1 / temperature，默认值0.07
 1. learnable_temperature: 温度是否可学习（训练中会被clamp防止溢出），默认值true
 
-该损失通常与逐点损失（如binary_cross_entropy）搭配使用，用同级的`weight`调节相对权重，经验值0.1量级
+注意事项：
+
+1. 一个list需要同时含有至少一个正样本和一个负样本才会计入该项损失，被屏蔽的list仍计入分母。使用`session_name`时要求同一个session的样本落在同一个batch内（不要求相邻），对于非DlrmHSTU系列模型样本构造方式与[jrc_loss](#jrc_loss)相同（`DISTRIBUTE BY session_id SORT BY session_id`），并尽量加大`batch_size`；若样本未按session分组，绝大多数list只含一条样本而被屏蔽，表现为该项损失接近0
+1. 该损失只约束list内的相对打分，不保证概率校准，通常与逐点损失（如binary_cross_entropy）搭配使用，用同级的`weight`调节相对权重，经验值0.1量级
 
 ## pe_mtl_loss
 

tzrec/loss/jrc_loss.py

@@ -10,27 +10,33 @@
 # limitations under the License.
 
 
+from typing import Optional
+
 import torch
+import torch.nn.functional as F
 from torch import Tensor
 from torch.nn import CrossEntropyLoss
 from torch.nn.modules.loss import _Loss
 
+from tzrec.ops.scatter_ops import lengths_to_index, scatter_logsumexp
 
-@torch.fx.wrap
-def _label_mask(labels: torch.Tensor) -> torch.Tensor:
-    return torch.eye(labels.size(0), dtype=torch.int64, device=labels.device)
-
-
-@torch.fx.wrap
-def _diag_index(labels: torch.Tensor) -> torch.Tensor:
-    return torch.arange(0, labels.size(0), dtype=torch.int64, device=labels.device)
+# Logit assigned to samples that must not compete: exp underflows to exactly
+# 0 after the segment max shift, so they drop out of the log-sum-exp.
+_MASKED_LOGIT = -1e9
 
 
 class JRCLoss(_Loss):
     """Positive sample probability competes in session.
 
     https://arxiv.org/abs/2208.06164
 
+    The session term of the paper is a softmax of each positive against the
+    negatives of its session (on the positive logit) and of each negative
+    against the positives of its session (on the negative logit).  Both
+    reduce to ``softplus(logsumexp(competitors) - own logit)``, where the
+    log-sum-exp is one value per session, so the term costs two segment
+    reductions instead of a batch-by-batch session mask.
+
     Args:
         alpha (float): cross entropy loss weight.
         reduction (str, optional): Specifies the reduction to apply to the
@@ -52,69 +58,50 @@ def forward(
         self,
         logits: Tensor,
         labels: Tensor,
-        session_ids: Tensor,
+        lengths: Tensor,
+        index: Optional[Tensor] = None,
     ) -> Tensor:
         """JRC loss.
 
+        Without ``index``, ``logits`` and ``labels`` are laid out
+        session by session in ``lengths`` order, each session contiguous.
+
         Args:
             logits: a `Tensor` with shape [batch_size, 2].
             labels: a `Tensor` with shape [batch_size].
-            session_ids: a `Tensor` with shape [batch_size].
+            lengths: a `Tensor` with shape [num_sessions], samples per
+                session, summing to batch_size.
+            index: a `Tensor` with shape [batch_size], session of each
+                sample (torch_scatter's ``index``), for samples in arbitrary
+                order.
 
         Return:
             loss: a `Tensor` with shape [batch_size] if reduction is 'none',
                     otherwise with shape ().
         """
         ce_loss = self._ce_loss(logits, labels)
 
-        batch_size = labels.shape[0]
-        mask = torch.eq(session_ids.unsqueeze(1), session_ids.unsqueeze(0)).float()
-        diag_index = _diag_index(labels)
+        if index is None:
+            index = lengths_to_index(lengths, output_size=logits.size(0))
         logits_neg, logits_pos = logits[:, 0], logits[:, 1]
-        diag = _label_mask(labels)
-        pos_num = torch.sum(labels)
-        neg_num = batch_size - pos_num
-
-        # first, we calculate pos sample loss in during the session.
-        pos_mask_index = torch.where(labels == 1.0)[0]
-        pos_diag_label = torch.index_select(diag_index, 0, pos_mask_index)
-        # pyre-ignore [6]
-        logits_pos = logits_pos.unsqueeze(0).tile([pos_num, 1])
-        pos_session_mask = torch.index_select(mask, 0, pos_mask_index)
-        # pyre-ignore [6]
-        y_pos = labels.unsqueeze(0).tile([pos_num, 1])
-        diag_pos = torch.index_select(diag, 0, pos_mask_index)
-        # we mask not in the same session, is diagonal and is positive.
-        logits_pos = (
-            logits_pos + ((1 - pos_session_mask) + (1 - diag_pos) * y_pos) * -1e9
-        )
-        loss_pos = self._ce_loss(logits_pos, pos_diag_label)
-
-        # next, we calculate neg sample loss in during the session.
-        neg_mask_index = torch.where(labels == 0.0)[0]
-        neg_diag_label = torch.index_select(diag_index, 0, neg_mask_index)
-        # neg_num is a 0-d integer tensor, which tile accepts via __index__.
-        # pyrefly: ignore[no-matching-overload]
-        logits_neg = logits_neg.unsqueeze(0).tile([neg_num, 1])
-        neg_session_mask = torch.index_select(mask, 0, neg_mask_index)
-        # pyrefly: ignore[no-matching-overload]
-        y_neg = (1 - labels).unsqueeze(0).tile([neg_num, 1])
-        diag_neg = torch.index_select(diag, 0, neg_mask_index)
-        # we mask not in the same session, is diagonal and is negative.
-        logits_neg = (
-            logits_neg + ((1 - neg_session_mask) + (1 - diag_neg) * y_neg) * -1e9
+        is_pos = labels == 1
+        # Competitors of a positive: the session's negatives on the positive
+        # logit; of a negative: the session's positives on the negative logit.
+        neg_competitors = torch.where(is_pos, _MASKED_LOGIT, logits_pos)
+        pos_competitors = torch.where(is_pos, logits_neg, _MASKED_LOGIT)
+        num_sessions = lengths.size(0)
+        lse_neg = scatter_logsumexp(
+            neg_competitors.unsqueeze(-1), index, num_sessions
+        ).squeeze(-1)
+        lse_pos = scatter_logsumexp(
+            pos_competitors.unsqueeze(-1), index, num_sessions
+        ).squeeze(-1)
+        ge_loss = torch.where(
+            is_pos,
+            F.softplus(lse_neg.index_select(0, index) - logits_pos),
+            F.softplus(lse_pos.index_select(0, index) - logits_neg),
         )
-        loss_neg = self._ce_loss(logits_neg, neg_diag_label)
-
         if self._reduction != "none":
-            loss_pos = loss_pos * pos_num / batch_size
-            loss_neg = loss_neg * neg_num / batch_size
-            ge_loss = loss_pos + loss_neg
-        else:
-            ge_loss = torch.zeros_like(labels, dtype=torch.float)
-            ge_loss.index_put_(torch.where(labels == 1.0), loss_pos)
-            ge_loss.index_put_(torch.where(labels == 0.0), loss_neg)
-
-        loss = self._alpha * ce_loss + (1 - self._alpha) * ge_loss
-        # pyre-ignore [7]
-        return loss
+            ge_loss = ge_loss.mean()
+
+        return self._alpha * ce_loss + (1 - self._alpha) * ge_loss

tzrec/loss/jrc_loss_test.py

@@ -13,53 +13,114 @@
 import unittest
 
 import torch
+from parameterized import parameterized
 
 from tzrec.loss.jrc_loss import JRCLoss
+from tzrec.utils.test_util import parameterized_name_func
+
+_LOGITS = torch.tensor(
+    [
+        [0.9, 0.1],
+        [0.5, 0.5],
+        [0.3, 0.7],
+        [0.2, 0.8],
+        [0.8, 0.2],
+        [0.55, 0.45],
+        [0.33, 0.67],
+        [0.55, 0.45],
+    ],
+    dtype=torch.float32,
+)
+_LABELS = torch.tensor([0, 0, 1, 1, 0, 0, 1, 1])
+_LENGTHS = torch.tensor([4, 4])
+
+
+def _reference_session_loss(
+    logits: torch.Tensor, labels: torch.Tensor, session_ids: torch.Tensor
+) -> torch.Tensor:
+    """Per-sample session term as the original batch-mask formulation.
+
+    Each positive competes with the negatives of its session on the positive
+    logit, each negative with the positives of its session on the negative
+    logit, the sample itself being the softmax target.
+    """
+    mask = session_ids.unsqueeze(1) == session_ids.unsqueeze(0)
+    is_pos = labels == 1
+    same_and_other_class = mask & (is_pos.unsqueeze(1) != is_pos.unsqueeze(0))
+    keep = same_and_other_class | torch.eye(labels.numel(), dtype=torch.bool)
+    own_channel = torch.where(is_pos, 1, 0)
+    rows = logits[:, own_channel].T
+    rows = rows.masked_fill(~keep, -1e9)
+    return torch.nn.functional.cross_entropy(
+        rows, torch.arange(labels.numel()), reduction="none"
+    )
 
 
 class JRCLossTest(unittest.TestCase):
     def test_jrc_loss(self) -> None:
-        loss_class = JRCLoss()
-        logits = torch.tensor(
-            [
-                [0.9, 0.1],
-                [0.5, 0.5],
-                [0.3, 0.7],
-                [0.2, 0.8],
-                [0.8, 0.2],
-                [0.55, 0.45],
-                [0.33, 0.67],
-                [0.55, 0.45],
-            ],
-            dtype=torch.float32,
-        )
-        labels = torch.tensor([0, 0, 1, 1, 0, 0, 1, 1])
-        session_ids = torch.tensor([1, 1, 1, 1, 2, 2, 2, 2], dtype=torch.int8)
-        loss = loss_class(logits, labels, session_ids)
+        loss = JRCLoss()(_LOGITS, _LABELS, _LENGTHS)
         self.assertEqual(0.7199, round(loss.item(), 4))
 
-
-class JRCLossTestReduceNone(unittest.TestCase):
     def test_jrc_loss_reduce_none(self) -> None:
-        loss_class = JRCLoss(reduction="none")
-        logits = torch.tensor(
+        loss = JRCLoss(reduction="none")(_LOGITS, _LABELS, _LENGTHS)
+        self.assertEqual((8,), tuple(loss.shape))
+        self.assertEqual(0.7199, round(torch.mean(loss).item(), 4))
+
+    @parameterized.expand(
+        [
+            # sessions with both classes, all positives, all negatives, one row
             [
-                [0.9, 0.1],
-                [0.5, 0.5],
-                [0.3, 0.7],
-                [0.2, 0.8],
-                [0.8, 0.2],
-                [0.55, 0.45],
-                [0.33, 0.67],
-                [0.55, 0.45],
+                [5, 3, 4, 1, 6],
+                [1, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 0, 1, 0, 1, 1, 0, 0, 1],
             ],
-            dtype=torch.float32,
+            [[19], [0] * 18 + [1]],
+            [[7, 7, 5], [1] * 7 + [0] * 7 + [1, 0, 1, 0, 1]],
+        ],
+        name_func=parameterized_name_func,
+    )
+    def test_matches_batch_mask_formulation(self, lengths, labels) -> None:
+        torch.manual_seed(0)
+        lengths_t = torch.tensor(lengths)
+        labels_t = torch.tensor(labels)
+        logits = torch.randn(labels_t.numel(), 2, requires_grad=True)
+        session_ids = torch.repeat_interleave(torch.arange(len(lengths)), lengths_t)
+
+        loss = JRCLoss(alpha=0.3, reduction="none")(logits, labels_t, lengths_t)
+        ce = torch.nn.functional.cross_entropy(logits, labels_t, reduction="none")
+        expected = 0.3 * ce + 0.7 * _reference_session_loss(
+            logits, labels_t, session_ids
         )
-        labels = torch.tensor([0, 0, 1, 1, 0, 0, 1, 1])
-        session_ids = torch.tensor([1, 1, 1, 1, 2, 2, 2, 2], dtype=torch.int8)
-        loss = loss_class(logits, labels, session_ids)
+        torch.testing.assert_close(loss, expected)
 
-        self.assertEqual(0.7199, round(torch.mean(loss).item(), 4))
+        (grad,) = torch.autograd.grad(loss.sum(), logits)
+        (expected_grad,) = torch.autograd.grad(expected.sum(), logits)
+        torch.testing.assert_close(grad, expected_grad)
+
+    def test_index_in_any_order(self) -> None:
+        torch.manual_seed(0)
+        lengths = torch.tensor([5, 3, 4])
+        labels = torch.tensor([1, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1])
+        logits = torch.randn(labels.numel(), 2)
+        perm = torch.randperm(labels.numel())
+        index = torch.repeat_interleave(torch.arange(3), lengths)[perm]
+
+        loss = JRCLoss(reduction="none")(logits[perm], labels[perm], lengths, index)
+        expected = JRCLoss(reduction="none")(logits, labels, lengths)[perm]
+        torch.testing.assert_close(loss, expected)
+
+    def test_bf16_logits_track_fp32(self) -> None:
+        """bf16 keeps the masked logit finite, incl. one-class sessions."""
+        torch.manual_seed(0)
+        lengths = torch.tensor([5, 3, 4, 1])
+        labels = torch.tensor([1, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 0, 1])
+        logits = torch.randn(labels.numel(), 2)
+        loss_fn = JRCLoss(reduction="none")
+        loss = loss_fn(logits.bfloat16(), labels, lengths)
+        self.assertEqual(loss.dtype, torch.bfloat16)
+        self.assertTrue(torch.isfinite(loss).all())
+        torch.testing.assert_close(
+            loss.float(), loss_fn(logits, labels, lengths), atol=5e-2, rtol=2e-2
+        )
 
 
 if __name__ == "__main__":