Add hard-negative token ranking objective

modeling_betterwright.py

@@ -21,6 +21,8 @@ class BetterWrightEncoderOutput(ModelOutput):
     logits: Optional[torch.Tensor] = None
     token_logits: Optional[torch.Tensor] = None
     uncertainty_logits: Optional[torch.Tensor] = None
+    token_loss: Optional[torch.Tensor] = None
+    token_rank_loss: Optional[torch.Tensor] = None
     last_hidden_state: Optional[torch.Tensor] = None
 
 
@@ -84,6 +86,8 @@ class BetterWrightEncoder(Lfm2PreTrainedModel):
         logits = self.relevance_head(pooled).squeeze(-1)
         uncertainty_logits = self.uncertainty_head(pooled).squeeze(-1)
         loss = None
+        token_loss = None
+        token_rank_loss = None
         if labels is not None:
             labels = labels.to(logits.dtype)
             positive_weight = float(getattr(self.config, "betterwright_positive_weight", 3.0))
@@ -122,12 +126,41 @@ class BetterWrightEncoder(Lfm2PreTrainedModel):
                         dtype=token_logits.dtype,
                     ),
                 )
-                loss = token_loss if loss is None else loss + 0.35 * token_loss
+                # Runtime pruning ranks lines by their highest-scoring token.
+                # BCE alone can achieve a reasonable average while leaving a
+                # few distractor tokens above the evidence. Explicitly compare
+                # each positive row against its hardest negative tokens so the
+                # training objective matches that production ranking behavior.
+                token_rank_losses = []
+                for row_index in range(token_logits.shape[0]):
+                    row_valid = valid[row_index]
+                    row_labels = token_labels[row_index]
+                    positives = token_logits[row_index][row_valid & (row_labels == 1)]
+                    negatives = token_logits[row_index][row_valid & (row_labels == 0)]
+                    if positives.numel() == 0 or negatives.numel() == 0:
+                        continue
+                    hard_count = min(64, negatives.numel())
+                    hard_negatives = torch.topk(
+                        negatives.float(), hard_count, sorted=False
+                    ).values
+                    token_rank_losses.append(
+                        F.softplus(
+                            1.0 + hard_negatives.mean() - positives.float().mean()
+                        )
+                    )
+                if token_rank_losses:
+                    token_rank_loss = torch.stack(token_rank_losses).mean()
+                else:
+                    token_rank_loss = token_loss.new_zeros(())
+                token_objective = 0.75 * token_loss + token_rank_loss
+                loss = token_objective if loss is None else loss + token_objective
         result = BetterWrightEncoderOutput(
             loss=loss,
             logits=logits,
             token_logits=token_logits,
             uncertainty_logits=uncertainty_logits,
+            token_loss=token_loss,
+            token_rank_loss=token_rank_loss,
             last_hidden_state=hidden,
         )
         return result if return_dict else tuple(result.values())