Add negative log-acceptance (LK) loss for speculative decoding (#657)

## Summary
Adds a negative log-acceptance loss to the speculators loss interface —
the
hyperparameter-free LK loss variant.

## Motivation
Acceptance rate equals the target/draft overlap, `alpha = sum_v min(p_v,
q_v) = 1 - TV`.
This loss is `-log(alpha)`, which optimizes acceptance directly. Its
gradient is
`(1/alpha)*grad(TV)`: the `1/alpha` factor amplifies TV's
otherwise-vanishing gradient when
overlap is low, so it trains well from a cold start where pure TV
stalls. At a point-mass
target it reduces to cross-entropy. Source: Samarin et al., "LK Losses:
Direct Acceptance
Rate Optimization for Speculative Decoding" (arXiv 2602.23881).

## Changes
- `neg_log_acceptance_loss(logits, targets) -> [1, seq_len]` in
`metrics.py`, reusing the
  same overlap term as `tv_loss`, with an `_EPS` floor before the log.
- Registered as `"nla"` in `resolve_loss_fn` (short name open to
maintainer preference).
- Unit tests: identical-dist = 0, equals `-log(overlap)`, reduces to
cross-entropy at a
  point-mass target, finite at zero overlap, and resolution.

## Scope / notes
- Negative-log-acceptance variant only. The hybrid LK loss (`lambda*KL +
(1-lambda)*TV` with an
adaptive, eta-controlled weight) is deferred to a follow-up, since it
needs a
loss-hyperparameter interface decision. The e2e variant and kernel
optimizations are out of
  scope.
- Target distribution is treated as the (detached) reference, consistent
with the existing
  `kl_div`/`ce`/`tv` losses which assume frozen-target logits.

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

Author: GIREESH7963

src/speculators/models/metrics.py

@@ -146,6 +146,35 @@ def tv_loss(
     return elementwise_loss  # noqa: RET504
 
 
+def neg_log_acceptance_loss(
+    logits: torch.Tensor,  # shape: [1, seq_len, draft_vocab_size]
+    targets: torch.Tensor,  # shape: [1, seq_len, draft_vocab_size]
+):
+    """Compute per-position negative log-acceptance (LK) loss.
+
+    The speculative-decoding acceptance rate equals the draft/target distribution
+    overlap, ``alpha = sum_v min(p_v, q_v)`` (the same quantity computed in
+    ``tv_loss``). This loss is ``-log(alpha)``. Its gradient is
+    ``(1 / alpha) * grad(TV)``: the ``1 / alpha`` factor amplifies the otherwise
+    vanishing TV gradient when overlap is low (early training), giving TV's
+    acceptance-optimal target a usable gradient from a cold start. When the target
+    is a point mass, this loss reduces to cross-entropy.
+
+    Args:
+        logits: Draft model logits (softmax applied internally to form q).
+        targets: Target model logits (softmax applied internally to form p).
+
+    Returns:
+        Per-position negative log-acceptance 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]
+    elementwise_loss = -torch.log(overlap.clamp_min(_EPS))
+
+    return elementwise_loss  # noqa: RET504
+
+
 def dflash_loss_decay(pos_idx: torch.Tensor, gamma: float):
     """Compute DFlash-style exponential decay weights per position.
 
@@ -186,8 +215,9 @@ def resolve_loss_fn(
     """Resolves a loss function given its abbreviated name.
 
     Args:
-        name: ``"kl_div"`` for KL-divergence, ``"ce"`` for cross-entropy, or
-            ``"tv"`` for total variation.
+        name: ``"kl_div"`` for KL-divergence, ``"ce"`` for cross-entropy,
+            ``"tv"`` for total variation, or ``"nla"`` for negative
+            log-acceptance.
 
     Returns:
         The corresponding loss function.
@@ -199,6 +229,7 @@ def resolve_loss_fn(
         "kl_div": kl_div_loss,
         "ce": ce_loss,
         "tv": tv_loss,
+        "nla": neg_log_acceptance_loss,
     }
     if name not in loss_fn_map:
         raise ValueError(

tests/unit/models/test_metrics.py

@@ -6,11 +6,13 @@
 import torch
 
 from speculators.models.metrics import (
+    ce_loss,
     compute_accuracy_single_step,
     dflash_loss_decay,
     exp_loss_decay,
     kl_div_loss,
     loss_function,
+    neg_log_acceptance_loss,
     resolve_loss_fn,
     tv_loss,
 )
@@ -118,6 +120,50 @@ def test_resolve_tv(self):
         assert resolve_loss_fn("tv") is tv_loss
 
 
+class TestNegLogAcceptanceLoss:
+    def test_identical_is_zero(self):
+        """Negative log-acceptance of identical distributions is ~0."""
+        logits = torch.randn(1, 4, 50)
+        loss = neg_log_acceptance_loss(logits, logits)
+        assert torch.allclose(loss, torch.zeros(1, 4), atol=1e-5)
+
+    def test_equals_neg_log_overlap_and_shape(self):
+        """Loss equals -log(overlap); output is [1, seq_len] and non-negative."""
+        torch.manual_seed(0)
+        logits = torch.randn(1, 4, 50)
+        targets = torch.randn(1, 4, 50)
+        out = neg_log_acceptance_loss(logits, targets)
+
+        overlap = 1.0 - tv_loss(logits, targets)  # tv = 1 - alpha
+        assert out.shape == (1, 4)
+        assert torch.allclose(out, -torch.log(overlap.clamp_min(1e-5)), atol=1e-6)
+        assert (out >= 0).all()
+
+    def test_reduces_to_cross_entropy_at_point_mass_target(self):
+        """At a point-mass target the loss reduces to cross-entropy."""
+        torch.manual_seed(0)
+        logits = torch.randn(1, 1, 50)
+        targets = torch.full((1, 1, 50), -30.0)
+        targets[0, 0, 7] = 30.0  # ~one-hot target
+        assert torch.allclose(
+            neg_log_acceptance_loss(logits, targets),
+            ce_loss(logits, targets),
+            atol=1e-3,
+        )
+
+    def test_finite_at_zero_overlap(self):
+        """The _EPS floor keeps the loss finite when overlap collapses to ~0."""
+        logits = torch.full((1, 1, 50), -30.0)
+        logits[0, 0, 0] = 30.0
+        targets = torch.full((1, 1, 50), -30.0)
+        targets[0, 0, 1] = 30.0
+        assert torch.isfinite(neg_log_acceptance_loss(logits, targets)).all()
+
+    def test_resolve_nla(self):
+        """resolve_loss_fn maps 'nla' to neg_log_acceptance_loss."""
+        assert resolve_loss_fn("nla") is neg_log_acceptance_loss
+
+
 class TestComputeAccuracySingleStep:
     def test_prev_correct_chain(self):
         """Conditional accuracy across ttt steps tracks cumulative correctness.