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

## Summary
Adds the hybrid LK loss: an adaptive per-position blend of KL and TV,
L = lambda*KL + (1-lambda)*TV with lambda = exp(-eta * sg[alpha]), alpha
= overlap.

## Motivation
TV optimizes acceptance directly but has a vanishing gradient from a
cold start; KL has a
strong early gradient but optimizes a proxy. The hybrid leans on KL when
overlap is low and
shifts to TV as the draft aligns, making TV's acceptance-optimal target
trainable from
scratch. Source: Samarin et al., arXiv 2602.23881. Paper's best setting:
eta=3.

## Changes
- `lk_hybrid_loss(logits, targets, eta=3.0)` in metrics.py, reusing
`kl_div_loss` and the
  TV overlap term; alpha in the blend weight is detached.
- `resolve_loss_fn` gains `lk_loss_eta` (default 3.0), binding eta into
the "lk_hybrid"
entry via functools.partial (single binding point; call sites
unchanged).
- `--lk-loss-eta` CLI flag (mirrors --dflash-decay-gamma), threaded
through each
get_trainer_kwargs into resolve_loss_fn; "lk_hybrid" added to --loss-fn
choices.
- Tests: eta->0 == KL, eta->inf == TV, shape/finiteness, alpha-detach
gradient check,
  resolve binds eta.

## Notes
- The adaptive KL/TV blend and the DFlash positional decay are
orthogonal weightings.
- tv/nla are not in the --loss-fn choices list either; happy to add them
here if desired.
- Loss key name "lk_hybrid" open to preference.

Signed-off-by: GIREESH7963 <abburigireesh@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: GIREESH7963

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.