Add reverse KL divergence (RKL) loss for speculative decoding

Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>

Author: guan404ming

scripts/train.py

@@ -864,7 +864,7 @@ def parse_args():
         default="kl_div",
         help=(
             "Loss function specification. Pass a name for a single loss "
-            "(kl_div, ce, tv, nla, lk_hybrid) or a JSON dict for a weighted "
+            "(kl_div, rkl, ce, tv, nla, lk_hybrid) or a JSON dict for a weighted "
             'combination, e.g. \'{"ce": 0.1, "tv": 0.9}\'.'
         ),
     )

src/speculators/models/metrics.py

@@ -98,6 +98,28 @@ def kl_div_loss(
     return elementwise_loss  # noqa: RET504
 
 
+def reverse_kl_div_loss(
+    logits: torch.Tensor,  # shape: [1, seq_len, draft_vocab_size]
+    targets: torch.Tensor,  # shape: [1, seq_len, draft_vocab_size]
+):
+    """Compute per-position reverse KL divergence from draft logits to target logits.
+
+    Args:
+        logits: Draft model logits (log-softmax applied internally).
+        targets: Target model logits (log-softmax applied internally).
+
+    Returns:
+        Per-position reverse KL divergence with shape [1, seq_len].
+    """
+    draft_logq = torch.nn.functional.log_softmax(logits, dim=-1)
+    target_logp = torch.nn.functional.log_softmax(targets, dim=-1)
+    elementwise_loss = torch.nn.functional.kl_div(
+        target_logp, draft_logq, reduction="none", log_target=True
+    ).sum(dim=-1)  # shape: [1, seq_len]
+
+    return elementwise_loss  # noqa: RET504
+
+
 def ce_loss(
     logits: torch.Tensor,  # shape: [1, seq_len, draft_vocab_size]
     targets: torch.Tensor,  # shape: [1, seq_len, draft_vocab_size]
@@ -258,6 +280,7 @@ def exp_loss_decay(pos_idx: torch.Tensor, gamma: float):
 
 _LOSS_FN_MAP: dict[str, Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = {
     "kl_div": kl_div_loss,
+    "rkl": reverse_kl_div_loss,
     "ce": ce_loss,
     "tv": tv_loss,
     "nla": neg_log_acceptance_loss,
@@ -271,9 +294,10 @@ def resolve_loss_fn(
     """Resolves a loss function given its abbreviated name.
 
     Args:
-        name: ``"kl_div"`` for KL-divergence, ``"ce"`` for cross-entropy,
-            ``"tv"`` for total variation, ``"nla"`` for negative
-            log-acceptance, or ``"lk_hybrid"`` for the adaptive KL/TV blend.
+        name: ``"kl_div"`` for KL-divergence, ``"rkl"`` for reverse KL-divergence,
+            ``"ce"`` for cross-entropy, ``"tv"`` for total variation, ``"nla"``
+            for negative log-acceptance, or ``"lk_hybrid"`` for the adaptive
+            KL/TV blend.
 
     Returns:
         The corresponding loss function.

tests/unit/models/test_metrics.py

@@ -15,6 +15,7 @@
     loss_function,
     neg_log_acceptance_loss,
     resolve_loss_fn,
+    reverse_kl_div_loss,
     tv_loss,
 )
 
@@ -93,6 +94,33 @@ def test_identical_zero_and_random_nonnegative(self):
         assert (loss_random >= -1e-6).all()
 
 
+class TestReverseKLDivLoss:
+    def test_identical_zero_and_random_nonnegative(self):
+        """Reverse KL of identical distributions is ~0; random inputs are >= 0."""
+        x = torch.randn(1, 4, 8)
+        loss_identical = reverse_kl_div_loss(x, x)
+        assert loss_identical.sum().item() == pytest.approx(0.0, abs=1e-4)
+
+        torch.manual_seed(0)
+        loss_random = reverse_kl_div_loss(torch.randn(1, 4, 8), torch.randn(1, 4, 8))
+        assert (loss_random >= -1e-6).all()
+
+    def test_equals_forward_kl_with_swapped_args(self):
+        """KL(q||p) equals forward KL with draft and target swapped."""
+        torch.manual_seed(0)
+        logits = torch.randn(1, 4, 8)
+        targets = torch.randn(1, 4, 8)
+        assert torch.allclose(
+            reverse_kl_div_loss(logits, targets),
+            kl_div_loss(targets, logits),
+            atol=1e-6,
+        )
+
+    def test_resolve_rkl(self):
+        """resolve_loss_fn maps 'rkl' to reverse_kl_div_loss."""
+        assert resolve_loss_fn("rkl") is reverse_kl_div_loss
+
+
 class TestTVLoss:
     def test_identical_is_zero(self):
         """TV distance between identical distributions is ~0."""