Add native loss support for JEV decision models

docs/source/en/model_doc/jev.md

@@ -81,6 +81,21 @@ and [d3](https://huggingface.co/vllm-sr/d3), whose Qwen3.5 backbone attends to t
 pipeline("decision", model="vllm-sr/d3")(request)  # or perplexity-ai/pplx-decider-v1.1-27b
 ```
 
+## Training
+
+`JevForDecision` returns `loss` and `logits` and can train with the standard `Trainer` or TRL.
+The default loss follows [AutoTrust JEV-9B](https://huggingface.co/autotrust/JEV-9B#training-details): KL divergence
+over each question's active options, plus `0.5 × RPS` for ordered Score questions. Set `config.rps_weight` to adjust
+the ordinal term. Training uses raw logits; calibration is applied by `decide`.
+
+Pass integer `labels` of shape `(1, num_questions)`, using `-100` to ignore a question, or soft target distributions
+of shape `(1, total_num_options)`, concatenated in question order. Each distribution must sum to one.
+`outputs.logits` uses the same packed option order.
+
+Optional `loss_weights` weight individual questions. The loss is averaged over valid questions, rather than the sum
+of weights. When accumulating prompts with different question counts, pass their total valid-question count as
+`num_items_in_batch` explicitly.
+
 ## JevConfig
 
 [[autodoc]] JevConfig

src/transformers/decision/modeling_utils.py

@@ -111,15 +111,15 @@ def decide(
         `Question.answer` turns into answers.
         """
         decision_config = decision_config or self.decision_config
-        inputs = {"input_ids": input_ids, **kwargs}
+        inputs = {"input_ids": input_ids, **kwargs, "return_dict": True}
         num_options = torch.tensor([[len(question.options) for question in decision_questions]])
         question_types = [Question.questions_types.index(question.type) for question in decision_questions]
         scores = self(
             marks=marks[None].to(self.device),
             num_options=num_options.to(self.device),
             question_type_ids=torch.tensor([question_types], device=self.device),
             **{key: value.to(self.device) if hasattr(value, "to") else value for key, value in inputs.items()},
-        )
+        ).logits
         return process_scores(
             decision_config, inputs["input_ids"], marks, decision_questions, scores.split(num_options[0].tolist())
         )

src/transformers/loss/loss_jev.py

@@ -0,0 +1,72 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+from torch.nn import functional as F
+
+
+def JevForDecisionLoss(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    num_options: torch.Tensor,
+    is_score: torch.Tensor,
+    rps_weight: float = 0.5,
+    loss_weights: torch.Tensor | None = None,
+    num_items_in_batch: torch.Tensor | int | None = None,
+    ignore_index: int = -100,
+    **kwargs,
+) -> torch.Tensor:
+    """AutoTrust JEV's documented KL(target || model) + ordinal RPS objective over each question's active slots.
+
+    Source: autotrust/JEV-9B@7ac669b8fb476ba926491e3d1c39f9330d68dbed, README.md, "Training details".
+    """
+    logits, labels = logits.float(), labels.to(logits.device).reshape(-1)
+    num_options = num_options.to(logits.device).reshape(-1)
+    positions = torch.arange(num_options.max(), device=logits.device)
+    option_mask = positions[None] < num_options[:, None]
+    starts = num_options.cumsum(0) - num_options
+    indices = (starts[:, None] + positions).clamp_max(logits.numel() - 1)
+    scores = logits[indices].masked_fill(~option_mask, -torch.inf)
+    # Exclude the maximum before log1p to retain tiny losses and target gradients at confident predictions.
+    maximum = scores.argmax(-1, keepdim=True)
+    is_maximum = positions[None] == maximum
+    shifted = torch.where(is_maximum, 0.0, scores - scores.gather(-1, maximum))
+    others = shifted.exp().masked_fill(is_maximum, 0.0).sum(-1, keepdim=True)
+    log_probabilities = (shifted - torch.log1p(others)).masked_fill(~option_mask, 0.0)
+    if labels.is_floating_point():
+        if labels.numel() != logits.numel():
+            raise ValueError("Soft labels must contain one probability per active option.")
+        targets = labels[indices].float().clamp_min(0).masked_fill(~option_mask, 0.0)
+        valid = torch.ones_like(num_options, dtype=torch.bool)
+    else:
+        if labels.numel() != num_options.numel():
+            raise ValueError("Integer labels must contain one option index per question.")
+        valid = labels != ignore_index
+        if torch.any(valid & ((labels < 0) | (labels >= num_options))):
+            raise ValueError("Each integer label must select an active option of its question.")
+        targets = F.one_hot(labels.masked_fill(~valid, 0).long(), num_classes=positions.shape[0]).float()
+        targets = targets * valid[:, None]
+    # Keep target entropy: soft-target CE has the same gradients, but a different loss value.
+    losses = F.kl_div(log_probabilities, targets, reduction="none").sum(-1)
+    if rps_weight:
+        probabilities = log_probabilities.exp() * option_mask
+        difference = probabilities.cumsum(-1) - targets.cumsum(-1)
+        gaps = positions[None] < (num_options - 1)[:, None]
+        rps = (difference.square() * gaps).sum(-1) / (num_options - 1).clamp_min(1)
+        losses = losses + rps_weight * rps * is_score.to(logits.device).reshape(-1)
+    if loss_weights is not None:
+        losses = losses * loss_weights.to(logits.device).reshape(-1)
+    # Normalize over valid questions, including when question weights are supplied.
+    normalizer = valid.sum().clamp_min(1) if num_items_in_batch is None else num_items_in_batch
+    normalizer = torch.as_tensor(normalizer, device=logits.device)
+    return (losses * valid).sum() / normalizer

src/transformers/loss/loss_utils.py

@@ -23,6 +23,7 @@
 from .loss_deimv2 import Deimv2ForObjectDetectionLoss
 from .loss_for_object_detection import ForObjectDetectionLoss, ForSegmentationLoss
 from .loss_grounding_dino import GroundingDinoForObjectDetectionLoss
+from .loss_jev import JevForDecisionLoss
 from .loss_lsd import PocketTTSForConditionalGenerationLoss
 from .loss_lw_detr import LwDetrForObjectDetectionLoss
 from .loss_rf_detr import RfDetrForSegmentationLoss
@@ -172,6 +173,7 @@ def ForSemanticSegmentationLoss(
 
 
 LOSS_MAPPING = {
+    "JevForDecision": JevForDecisionLoss,
     "ForSemanticSegmentation": ForSemanticSegmentationLoss,
     "ForCausalLM": ForCausalLMLoss,
     "ForMaskedLM": ForMaskedLMLoss,

src/transformers/models/jev/configuration_jev.py

@@ -107,6 +107,8 @@ class JevConfig(PreTrainedConfig):
         `offsets[question.type]` on, in order.
     softcap (`float`, *optional*):
         Soft cap of the scores, `softcap * tanh(scores / softcap)`, like the logits of a language model head.
+    rps_weight (`float`, *optional*, defaults to 0.5):
+        Weight of the ordinal ranked probability score for `score` questions, added to the KL loss.
     text_config (`dict` or `PreTrainedConfig`, *optional*):
         Config of the model the head reads, the `config.json` of its repo.
 
@@ -124,6 +126,7 @@ class JevConfig(PreTrainedConfig):
     num_verbalizers: int = 24
     offsets: dict[str, int] | None = None
     softcap: float | None = None
+    rps_weight: float = 0.5
     text_config: dict | PreTrainedConfig | None = None
     initializer_range: float = 0.02
 

src/transformers/models/jev/modeling_jev.py

@@ -11,13 +11,15 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from dataclasses import dataclass
+
 import torch
 from torch import nn
 
 from ...decision import Question
 from ...decision.modeling_utils import DecisionMixin
 from ...modeling_utils import PreTrainedModel
-from ...utils import auto_docstring
+from ...utils import ModelOutput, auto_docstring, can_return_tuple
 from ..auto import AutoModel
 from .configuration_jev import JevConfig
 
@@ -29,28 +31,56 @@ class JevPreTrainedModel(PreTrainedModel):
     _supports_sdpa = _supports_flash_attn = _supports_flex_attn = True
 
 
+@auto_docstring
+@dataclass
+class JevDecisionOutput(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `()`, *optional*):
+        Mean KL loss plus the weighted ranked probability score on `score` questions, when labels are supplied.
+    logits (`torch.FloatTensor` of shape `(total_num_options,)`):
+        Uncalibrated active-option scores concatenated in question order.
+    """
+
+    loss: torch.Tensor | None = None
+    logits: torch.Tensor | None = None
+
+
 @auto_docstring(
     custom_intro="""
     JEV (autotrust/JEV-9B): the model of its repo and a decision head, linear scores of single-token labels at the token
     before the `<|mark|>` of each question, soft-capped like the logits of a language model head.
     """
 )
 class JevForDecision(DecisionMixin, JevPreTrainedModel):
+    # Packed soft labels count options, so Trainer must not infer a question count from their shape.
+    accepts_loss_kwargs = False
+
     def __init__(self, config: JevConfig):
         super().__init__(config)
         self.model = AutoModel.from_config(config.text_config)
         self.proj = nn.Linear(config.hidden_size, config.num_verbalizers)
         self.post_init()
 
+    @can_return_tuple
     @auto_docstring
     def forward(
         self,
         input_ids: torch.LongTensor,
         marks: torch.LongTensor,
         num_options: torch.LongTensor,
         question_type_ids: torch.LongTensor,
+        labels: torch.Tensor | None = None,
+        loss_weights: torch.Tensor | None = None,
         **kwargs,
-    ) -> torch.Tensor:
+    ) -> JevDecisionOutput:
+        r"""
+        labels (`torch.Tensor`, *optional*):
+            Integer option indices of shape `(1, num_questions)`, with `-100` for ignored questions, or floating-point
+            target distributions of shape `(1, total_num_options)`, concatenated in question order. Each distribution
+            must sum to one over that question's active options.
+        loss_weights (`torch.Tensor` of shape `(1, num_questions)`, *optional*):
+            Per-question weights. The weighted loss is divided by the number of valid questions, not the sum of weights.
+        """
         # A batch of one prompt
         marks, num_options, question_type_ids = marks[0], num_options[0], question_type_ids[0]
         hidden_states = self.model(input_ids=input_ids, use_cache=False, **kwargs).last_hidden_state[0]
@@ -61,7 +91,19 @@ def forward(
         offsets = [self.config.offsets[question_type] for question_type in Question.questions_types]
         starts = torch.tensor(offsets, device=logits.device)[question_type_ids]
         rows = torch.arange(logits.shape[-1], device=logits.device)
-        return logits[(rows >= starts[:, None]) & (rows < (starts + num_options)[:, None])]
+        scores = logits[(rows >= starts[:, None]) & (rows < (starts + num_options)[:, None])]
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=scores,
+                labels=labels,
+                num_options=num_options,
+                is_score=question_type_ids == Question.questions_types.index("score"),
+                rps_weight=self.config.rps_weight,
+                loss_weights=loss_weights,
+                **kwargs,
+            )
+        return JevDecisionOutput(loss=loss, logits=scores)
 
 
 __all__ = ["JevForDecision", "JevPreTrainedModel"]

tests/models/jev/test_modeling_jev.py

@@ -21,7 +21,7 @@
 from safetensors.torch import save_file
 
 from transformers import AutoProcessor, is_torch_available, pipeline
-from transformers.decision import DecisionConfig
+from transformers.decision import DecisionConfig, Question
 from transformers.testing_utils import require_torch, require_torch_accelerator, require_vision, slow
 
 from ...test_decision_head_common import DecisionHeadTesterMixin, check_decision, render
@@ -31,6 +31,89 @@
     import torch
 
     from transformers import JevForDecision
+    from transformers.loss.loss_jev import JevForDecisionLoss
+
+
+@require_torch
+class JevLossTest(unittest.TestCase):
+    def setUp(self):
+        self.logits = torch.tensor(
+            [1.0, 0.1, -0.2, -1.0, 0.5, 1.2, -0.4, 0.3, 0.8, 0.4, -0.5, 0.6, 0.9], requires_grad=True
+        )
+        self.num_options = torch.tensor([1, 2, 6, 4])
+        self.is_score = torch.tensor([True, False, True, False])
+        self.labels = torch.tensor([0, 1, 3, 2])
+        self.soft_labels = torch.tensor([1.0, 0.2, 0.8, 0.05, 0.1, 0.2, 0.3, 0.25, 0.1, 0.0, 0.2, 0.0, 0.8])
+        self.weights = torch.tensor([0.5, 0.25, 2.0, 0.5])
+
+    def test_model_card_loss_and_gradients(self):
+        # Independent float64 categorical KL and ordinal RPS reference, including ragged and singleton questions.
+        sizes = self.num_options.tolist()
+        ignored = self.labels.clone()
+        ignored[-1] = -100
+        for labels in (self.labels, self.soft_labels, ignored, torch.full_like(self.labels, -100)):
+            soft_labels = labels.is_floating_point()
+            for weights in (None, self.weights):
+                for rps_weight in (0.0, 0.5):
+                    with self.subTest(labels=labels.tolist(), weighted=weights is not None, rps_weight=rps_weight):
+                        reference_logits = self.logits.detach().double().requires_grad_()
+                        targets = labels.split(sizes) if soft_labels else labels
+                        terms = []
+                        for index, (scores, target) in enumerate(zip(reference_logits.split(sizes), targets)):
+                            if not soft_labels:
+                                if target == -100:
+                                    continue
+                                target = torch.nn.functional.one_hot(target, scores.numel())
+                            teacher = torch.distributions.Categorical(probs=target.double())
+                            student = torch.distributions.Categorical(logits=scores)
+                            term = torch.distributions.kl_divergence(teacher, student)
+                            if self.is_score[index] and scores.numel() > 1:
+                                difference = student.probs.cumsum(0)[:-1] - teacher.probs.cumsum(0)[:-1]
+                                term = term + rps_weight * difference.square().mean()
+                            terms.append(term if weights is None else term * weights[index])
+                        expected = torch.stack(terms).mean() if terms else reference_logits.sum() * 0
+                        actual = JevForDecisionLoss(
+                            self.logits,
+                            labels,
+                            self.num_options,
+                            self.is_score,
+                            rps_weight=rps_weight,
+                            loss_weights=weights,
+                        )
+                        torch.testing.assert_close(actual.double(), expected)
+                        actual_gradient = torch.autograd.grad(actual, self.logits)[0]
+                        expected_gradient = torch.autograd.grad(expected, reference_logits)[0]
+                        torch.testing.assert_close(actual_gradient.double(), expected_gradient)
+
+    def test_accumulated_loss(self):
+        loss = JevForDecisionLoss(
+            self.logits, self.soft_labels, self.num_options, self.is_score, loss_weights=self.weights
+        )
+        expected_gradient = torch.autograd.grad(loss, self.logits)[0]
+        start = 0
+        for index, count in enumerate(self.num_options.tolist()):
+            end = start + count
+            JevForDecisionLoss(
+                self.logits[start:end],
+                self.soft_labels[start:end],
+                self.num_options[index : index + 1],
+                self.is_score[index : index + 1],
+                loss_weights=self.weights[index : index + 1],
+                num_items_in_batch=4,
+            ).backward()
+            start = end
+        torch.testing.assert_close(self.logits.grad, expected_gradient)
+
+    def test_confident_prediction_gradients(self):
+        # Standard fp32 log_softmax rounds the winning loss and gradient to zero at this margin.
+        logits = torch.tensor([9.0, -9.0, -9.0], requires_grad=True)
+        loss = JevForDecisionLoss(logits, torch.tensor([0]), torch.tensor([3]), torch.tensor([False]))
+        mass = 2 * torch.exp(torch.tensor(-18.0, dtype=torch.float64))
+        non_target_gradient = mass / (1 + mass) / 2
+        expected_gradient = torch.stack([-2 * non_target_gradient, non_target_gradient, non_target_gradient])
+        loss.backward()
+        torch.testing.assert_close(loss.double(), torch.log1p(mass), rtol=1e-6, atol=0)
+        torch.testing.assert_close(logits.grad.double(), expected_gradient, rtol=1e-6, atol=0)
 
 
 @require_torch
@@ -39,6 +122,32 @@ class JevForDecisionTest(DecisionHeadTesterMixin, unittest.TestCase):
     decision_repo = "hf-internal-testing/decision-gev"
     tiny_sizes = {"hidden_size": 16}
 
+    def test_loss_backward(self):
+        inputs = {
+            "input_ids": torch.arange(12)[None],
+            "marks": torch.tensor([[4, 8, 12]]),
+            "num_options": torch.tensor([[2, 6, 4]]),
+            "question_type_ids": torch.tensor(
+                [[Question.questions_types.index(kind) for kind in ("noul", "score", "choice")]]
+            ),
+            "labels": torch.tensor([[1, 3, 2]]),
+        }
+        for dtype in (torch.float32, torch.bfloat16):
+            with self.subTest(dtype=dtype):
+                model = JevForDecision(self.config).to(dtype).eval()
+                output = model(**inputs)
+                self.assertEqual(output.loss.dtype, torch.float32)
+                output.loss.backward()
+                self.assertTrue(all(p.grad is None or torch.isfinite(p.grad).all() for p in model.parameters()))
+                self.assertGreater(model.proj.weight.grad.abs().sum(), 0)
+                self.assertGreater(model.model.get_input_embeddings().weight.grad.abs().sum(), 0)
+                active = torch.zeros(self.config.num_verbalizers, dtype=torch.bool)
+                for kind, count in zip(("noul", "score", "choice"), inputs["num_options"][0]):
+                    start = self.config.offsets[kind]
+                    active[start : start + count] = True
+                inactive_gradients = model.proj.weight.grad[~active]
+                torch.testing.assert_close(inactive_gradients, torch.zeros_like(inactive_gradients))
+
     @require_vision
     def test_chat_template_reproduces_the_model_code(self):
         super().test_chat_template_reproduces_the_model_code()