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

Signed-off-by: GIREESH7963 <abburigireesh@gmail.com>

Author: GIREESH7963

scripts/train.py

@@ -963,7 +963,7 @@ def parse_args():
         "--loss-fn",
         type=str,
         default="kl_div",
-        choices=["kl_div", "ce"],
+        choices=["kl_div", "ce", "lk_hybrid"],
         help=(
             "Loss function used during draft model training. "
             "'kl_div' = KL divergence (default). "
@@ -1047,6 +1047,13 @@ def parse_args():
         default=4.0,
         help="Decay gamma for DFlash loss weighting (default: 4.0)",
     )
+    parser.add_argument(
+        "--lk-loss-eta",
+        type=float,
+        default=3.0,
+        help="Blend temperature for the hybrid LK loss (--loss-fn lk_hybrid). "
+        "Higher values shift from KL toward TV sooner (default: 3.0).",
+    )
     parser.add_argument(
         "--draft-attn-impl",
         type=str,
@@ -1172,7 +1179,7 @@ def parse_args():
     args = parser.parse_args()
     provided = explicitly_provided_dests(parser, DECODER_SHAPING_FLAGS)
     validate_draft_init_args(parser, args, provided)
-    resolve_loss_fn(args.loss_fn)
+    resolve_loss_fn(args.loss_fn, lk_loss_eta=args.lk_loss_eta)
     return args
 
 

src/speculators/models/dflash/core.py

@@ -193,7 +193,9 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:
         Returns:
             Tuple of (train_call_kwargs, val_call_kwargs)
         """
-        loss_fn = resolve_loss_fn(kwargs["loss_fn"])
+        loss_fn = resolve_loss_fn(
+            kwargs["loss_fn"], lk_loss_eta=kwargs.get("lk_loss_eta", 3.0)
+        )
         gamma = kwargs.get("dflash_decay_gamma", 4.0)
         return {"loss_fn": loss_fn, "gamma": gamma}, {
             "loss_fn": loss_fn,

src/speculators/models/eagle3/core.py

@@ -361,7 +361,9 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:
         Returns:
             Tuple of (train_call_kwargs, val_call_kwargs)
         """
-        loss_fn = resolve_loss_fn(kwargs["loss_fn"])
+        loss_fn = resolve_loss_fn(
+            kwargs["loss_fn"], lk_loss_eta=kwargs.get("lk_loss_eta", 3.0)
+        )
         train_kwargs = {
             "use_off_policy_tokens": kwargs["use_off_policy_tokens"],
             "ttt_steps": kwargs["ttt_steps"],

src/speculators/models/metrics.py

@@ -1,4 +1,5 @@
 from collections.abc import Callable
+from functools import partial
 
 import torch
 
@@ -175,6 +176,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.
 
@@ -211,13 +254,17 @@ def exp_loss_decay(pos_idx: torch.Tensor, gamma: float):
 
 def resolve_loss_fn(
     name: str,
+    lk_loss_eta: float = 3.0,
 ) -> "Callable[[torch.Tensor, torch.Tensor], torch.Tensor]":
     """Resolves a loss function given its abbreviated name.
 
     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.
+        lk_loss_eta: Blend temperature for ``"lk_hybrid"``; bound into the loss
+            via ``functools.partial`` so callers keep handing a bare
+            ``(logits, targets)`` callable. Ignored by the other losses.
 
     Returns:
         The corresponding loss function.
@@ -230,6 +277,7 @@ def resolve_loss_fn(
         "ce": ce_loss,
         "tv": tv_loss,
         "nla": neg_log_acceptance_loss,
+        "lk_hybrid": partial(lk_hybrid_loss, eta=lk_loss_eta),
     }
     if name not in loss_fn_map:
         raise ValueError(

src/speculators/models/peagle/core.py

@@ -252,5 +252,7 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:
         Returns:
             Tuple of (train_call_kwargs, val_call_kwargs)
         """
-        loss_fn = resolve_loss_fn(kwargs["loss_fn"])
+        loss_fn = resolve_loss_fn(
+            kwargs["loss_fn"], lk_loss_eta=kwargs.get("lk_loss_eta", 3.0)
+        )
         return {"loss_fn": loss_fn}, {"loss_fn": loss_fn}

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,70 @@ 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_binds_eta(self):
+        """resolve_loss_fn binds eta into the 'lk_hybrid' callable via partial."""
+        torch.manual_seed(0)
+        logits, targets = torch.randn(1, 4, 50), torch.randn(1, 4, 50)
+        fn = resolve_loss_fn("lk_hybrid", lk_loss_eta=2.0)
+        assert torch.allclose(
+            fn(logits, targets), lk_hybrid_loss(logits, targets, eta=2.0), atol=1e-6
+        )
+
+
 class TestComputeAccuracySingleStep:
     def test_prev_correct_chain(self):
         """Conditional accuracy across ttt steps tracks cumulative correctness.