Add hybrid LK loss (adaptive KL/TV blend) for speculative decoding

scripts/train.py

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

src/speculators/models/metrics.py

@@ -180,6 +180,48 @@ def neg_log_acceptance_loss(
     return elementwise_loss  # noqa: RET504
 
 
+def lk_hybrid_loss(
+    logits: torch.Tensor,  # shape: [1, seq_len, draft_vocab_size]
+    targets: torch.Tensor,  # shape: [1, seq_len, draft_vocab_size]
+    eta: float = 3.0,
+):
+    """Compute per-position hybrid LK loss (adaptive KL/TV blend).
+
+    Blends KL divergence and total variation per position:
+    ``L = lambda * KL(p||q) + (1 - lambda) * TV(p, q)`` with adaptive weight
+    ``lambda = exp(-eta * sg[alpha])``, where ``alpha = sum_v min(p_v, q_v)`` is the
+    acceptance rate (overlap) and ``sg`` is stop-gradient. When overlap is low
+    (early training, misaligned draft) ``lambda -> 1`` and the loss leans on KL's
+    strong gradient; as overlap grows ``lambda -> 0`` and it shifts to TV, which
+    optimizes acceptance directly. This gives TV's acceptance-optimal target a
+    usable gradient from a cold start.
+
+    ``alpha`` in the weight is detached: it controls the blend but is not
+    differentiated through; gradients flow only through the KL and TV terms.
+
+    Source: Samarin et al., "LK Losses: Direct Acceptance Rate Optimization for
+    Speculative Decoding" (arXiv 2602.23881), hybrid objective.
+
+    Args:
+        logits: Draft model logits (softmax applied internally to form q).
+        targets: Target model logits (softmax applied internally to form p).
+        eta: Blend temperature; larger shifts toward TV sooner. Default 3.0
+            (the paper's best hybrid setting).
+
+    Returns:
+        Per-position hybrid loss with shape [1, seq_len].
+    """
+    draft_p = torch.nn.functional.softmax(logits, dim=-1)
+    target_p = torch.nn.functional.softmax(targets, dim=-1)
+    overlap = torch.minimum(draft_p, target_p).sum(dim=-1)  # alpha, shape: [1, seq_len]
+    tv = 1.0 - overlap
+    kl = kl_div_loss(logits, targets)  # reuse existing KL, shape: [1, seq_len]
+    weight = torch.exp(-eta * overlap.detach())  # lambda = exp(-eta * sg[alpha])
+    elementwise_loss = weight * kl + (1.0 - weight) * tv
+
+    return elementwise_loss  # noqa: RET504
+
+
 def dflash_loss_decay(pos_idx: torch.Tensor, gamma: float):
     """Compute DFlash-style exponential decay weights per position.
 
@@ -219,6 +261,7 @@ def exp_loss_decay(pos_idx: torch.Tensor, gamma: float):
     "ce": ce_loss,
     "tv": tv_loss,
     "nla": neg_log_acceptance_loss,
+    "lk_hybrid": lk_hybrid_loss,
 }
 
 
@@ -229,8 +272,8 @@ def resolve_loss_fn(
 
     Args:
         name: ``"kl_div"`` for KL-divergence, ``"ce"`` for cross-entropy,
-            ``"tv"`` for total variation, or ``"nla"`` for negative
-            log-acceptance.
+            ``"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

@@ -11,6 +11,7 @@
     dflash_loss_decay,
     exp_loss_decay,
     kl_div_loss,
+    lk_hybrid_loss,
     loss_function,
     neg_log_acceptance_loss,
     resolve_loss_fn,
@@ -164,6 +165,65 @@ def test_resolve_nla(self):
         assert resolve_loss_fn("nla") is neg_log_acceptance_loss
 
 
+class TestLKHybridLoss:
+    def test_eta_zero_reduces_to_kl(self):
+        """eta=0 gives lambda=1 everywhere, so the loss is pure KL."""
+        torch.manual_seed(0)
+        logits, targets = torch.randn(1, 4, 50), torch.randn(1, 4, 50)
+        assert torch.allclose(
+            lk_hybrid_loss(logits, targets, eta=0.0),
+            kl_div_loss(logits, targets),
+            atol=1e-6,
+        )
+
+    def test_large_eta_reduces_to_tv(self):
+        """Large eta drives lambda->0, so the loss approaches pure TV."""
+        torch.manual_seed(0)
+        logits, targets = torch.randn(1, 4, 50), torch.randn(1, 4, 50)
+        assert torch.allclose(
+            lk_hybrid_loss(logits, targets, eta=1e6),
+            tv_loss(logits, targets),
+            atol=1e-5,
+        )
+
+    def test_shape_and_finite(self):
+        """Output is [1, seq_len] and finite at the default blend setting."""
+        torch.manual_seed(0)
+        logits, targets = torch.randn(1, 4, 50), torch.randn(1, 4, 50)
+        out = lk_hybrid_loss(logits, targets, eta=3.0)
+        assert out.shape == (1, 4)
+        assert torch.isfinite(out).all()
+
+    def test_alpha_is_detached_in_weight(self):
+        """The alpha inside lambda must be stop-gradient.
+
+        Verify the impl's gradient matches the detached form, and that NOT
+        detaching would differ.
+        """
+        torch.manual_seed(0)
+        logits = torch.randn(1, 3, 40, requires_grad=True)
+        targets = torch.randn(1, 3, 40)
+        g_impl = torch.autograd.grad(
+            lk_hybrid_loss(logits, targets, eta=3.0).sum(), logits
+        )[0]
+
+        def manual(detach):
+            dp, tp = torch.softmax(logits, -1), torch.softmax(targets, -1)
+            ov = torch.minimum(dp, tp).sum(-1)
+            alpha = ov.detach() if detach else ov
+            lam = torch.exp(-3.0 * alpha)
+            return (lam * kl_div_loss(logits, targets) + (1 - lam) * (1 - ov)).sum()
+
+        g_detached = torch.autograd.grad(manual(True), logits, retain_graph=True)[0]
+        g_nodetach = torch.autograd.grad(manual(False), logits)[0]
+        assert torch.allclose(g_impl, g_detached, atol=1e-5)
+        assert not torch.allclose(g_detached, g_nodetach, atol=1e-4)
+
+    def test_resolve_lk_hybrid(self):
+        """resolve_loss_fn maps 'lk_hybrid' to lk_hybrid_loss."""
+        assert resolve_loss_fn("lk_hybrid") is lk_hybrid_loss
+
+
 class TestComputeAccuracySingleStep:
     def test_prev_correct_chain(self):
         """Conditional accuracy across ttt steps tracks cumulative correctness.