diff --git a/scripts/train.py b/scripts/train.py index e45349d20..bb1bf51d8 100644 --- a/scripts/train.py +++ b/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( diff --git a/src/speculators/models/metrics.py b/src/speculators/models/metrics.py index a1a6cb651..fc595f08f 100644 --- a/src/speculators/models/metrics.py +++ b/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. diff --git a/tests/unit/models/test_metrics.py b/tests/unit/models/test_metrics.py index 95aed0271..014f26003 100644 --- a/tests/unit/models/test_metrics.py +++ b/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.