[DFlash/DSpark] Optional direct-acceptance (LK) loss

README.md

@@ -38,6 +38,8 @@ bash scripts/train/train.sh
 
 Hardware: the default configs and scripts assume a single node with 8 GPUs. For fewer GPUs, reduce `CUDA_VISIBLE_DEVICES`.
 
+**Direct-acceptance (LK) loss (optional).** Set `accept_loss_alpha > 0` in a DFlash/DSpark config's `model` block to add a differentiable objective that maximises the drafter's expected accepted length E[tau] = sum_k prod_{j<=k} alpha_j directly, rather than relying only on the per-token cross-entropy proxy. It makes the acceptance quantity the trainer already logs as `tau_probabilistic` trainable; `0.0` (default) keeps prior behaviour.
+
 
 ## Evaluation
 

config/dflash/dflash_qwen3_8b.py

@@ -24,6 +24,11 @@
     loss_decay_gamma=4.0,
     ce_loss_alpha=1.0,
     l1_loss_alpha=0.0,
+
+    # Optional direct-acceptance (LK) loss. >0 trains the drafter to maximise
+    # expected accepted length E[tau]=sum_k prod_{j<=k} alpha_j directly,
+    # rather than only the per-token cross-entropy proxy. 0.0 keeps prior behaviour.
+    accept_loss_alpha=0.0,
 )
 
 train = dict(

deepspec/modeling/dspark/loss.py

@@ -14,7 +14,7 @@ def _all_reduce_loss_denominators(
     world_size: int,
 ) -> dict[str, torch.Tensor]:
     denominators = {}
-    for key in ("ce_loss_den", "l1_loss_den", "confidence_loss_den"):
+    for key in ("ce_loss_den", "l1_loss_den", "confidence_loss_den", "accept_loss_den"):
         tensor = loss_terms[key].detach().clone()
         if world_size > 1:
             dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
@@ -92,6 +92,7 @@ def _collect_local_terms(
     outputs: DSparkForwardOutput,
     loss_decay_gamma: Optional[float],
     l1_loss_alpha: float,
+    accept_loss_alpha: float,
 ) -> tuple[dict[str, torch.Tensor], bool]:
     draft_logits = outputs.draft_logits
     target_ids = outputs.target_ids
@@ -143,6 +144,21 @@ def _collect_local_terms(
             valid_block_weights=valid_block_weights,
         )
 
+    accept_loss_num = zero
+    accept_loss_den = zero
+    if accept_loss_alpha > 0:
+        assert accept_rate_3d is not None, (
+            "aligned_target_logits is required when accept_loss_alpha > 0."
+        )
+        # Trainable form of the ``tau_probabilistic`` metric: maximise the
+        # expected accepted draft length E[accepted] = sum_k prod_{j<=k} alpha_j,
+        # with alpha the per-position acceptance probability. Negated so that
+        # minimising the loss maximises accepted length.
+        valid_accept_rate = accept_rate_3d * eval_mask.to(torch.float32)
+        expected_draft_accepted = valid_accept_rate.cumprod(dim=-1).sum(dim=-1)
+        accept_loss_num = -(expected_draft_accepted * valid_block_weights).sum()
+        accept_loss_den = valid_block_weights.sum()
+
     has_confidence = outputs.confidence_pred is not None
     confidence_loss_num = zero
     confidence_loss_den = zero
@@ -187,6 +203,8 @@ def _collect_local_terms(
         "l1_loss_den": l1_loss_den,
         "confidence_loss_num": confidence_loss_num,
         "confidence_loss_den": confidence_loss_den,
+        "accept_loss_num": accept_loss_num,
+        "accept_loss_den": accept_loss_den,
     }
 
     for pos_idx in range(block_size):
@@ -231,6 +249,7 @@ def _build_loss(
     ce_loss_alpha: float,
     l1_loss_alpha: float,
     confidence_head_alpha: float,
+    accept_loss_alpha: float,
     has_confidence: bool,
     world_size: int,
 ) -> torch.Tensor:
@@ -245,10 +264,16 @@ def _build_loss(
         confidence_loss = loss_terms["confidence_loss_num"] / (
             global_denominators["confidence_loss_den"] + 1e-6
         )
+    accept_loss = ce_loss.new_zeros(())
+    if global_denominators["accept_loss_den"].item() > 0:
+        accept_loss = loss_terms["accept_loss_num"] / (
+            global_denominators["accept_loss_den"] + 1e-6
+        )
     return (
         ce_loss_alpha * ce_loss
         + l1_loss_alpha * l1_loss
         + confidence_head_alpha * confidence_loss
+        + accept_loss_alpha * accept_loss
     ) * world_size
 
 
@@ -259,11 +284,13 @@ def compute_dspark_loss(
     ce_loss_alpha: float,
     l1_loss_alpha: float,
     confidence_head_alpha: float,
+    accept_loss_alpha: float = 0.0,
 ):
     loss_terms, has_confidence = _collect_local_terms(
         outputs=outputs,
         loss_decay_gamma=loss_decay_gamma,
         l1_loss_alpha=float(l1_loss_alpha),
+        accept_loss_alpha=float(accept_loss_alpha),
     )
     world_size = dist.get_world_size()
     global_denominators = _all_reduce_loss_denominators(
@@ -273,6 +300,7 @@ def compute_dspark_loss(
     ce_loss_alpha = float(ce_loss_alpha)
     l1_loss_alpha = float(l1_loss_alpha)
     confidence_head_alpha = float(confidence_head_alpha)
+    accept_loss_alpha = float(accept_loss_alpha)
 
     local_ce_loss = loss_terms["ce_loss_num"] / (loss_terms["ce_loss_den"] + 1e-6)
     local_l1_loss = local_ce_loss.new_zeros(())
@@ -285,10 +313,16 @@ def compute_dspark_loss(
         local_confidence_loss = loss_terms["confidence_loss_num"] / (
             loss_terms["confidence_loss_den"] + 1e-6
         )
+    local_accept_loss = local_ce_loss.new_zeros(())
+    if loss_terms["accept_loss_den"].item() > 0:
+        local_accept_loss = loss_terms["accept_loss_num"] / (
+            loss_terms["accept_loss_den"] + 1e-6
+        )
     local_loss = (
         ce_loss_alpha * local_ce_loss
         + l1_loss_alpha * local_l1_loss
         + confidence_head_alpha * local_confidence_loss
+        + accept_loss_alpha * local_accept_loss
     )
 
     add_metric(
@@ -311,6 +345,13 @@ def compute_dspark_loss(
             den=loss_terms["confidence_loss_den"],
             tag="train",
         )
+    if loss_terms["accept_loss_den"].item() > 0:
+        add_metric(
+            "accept_loss",
+            loss_terms["accept_loss_num"],
+            den=loss_terms["accept_loss_den"],
+            tag="train",
+        )
     add_metric(
         "loss",
         local_loss,
@@ -323,6 +364,7 @@ def compute_dspark_loss(
         ce_loss_alpha=ce_loss_alpha,
         l1_loss_alpha=l1_loss_alpha,
         confidence_head_alpha=confidence_head_alpha,
+        accept_loss_alpha=accept_loss_alpha,
         has_confidence=has_confidence,
         world_size=world_size,
     )

deepspec/trainer/dspark_trainer.py

@@ -35,6 +35,7 @@ def run_batch(self, batch):
             ce_loss_alpha=float(self.args.model.ce_loss_alpha),
             l1_loss_alpha=float(self.args.model.l1_loss_alpha),
             confidence_head_alpha=float(self.args.model.confidence_head_alpha),
+            accept_loss_alpha=float(getattr(self.args.model, "accept_loss_alpha", 0.0)),
         )
         return loss
 

scripts/test_accept_loss.py

@@ -0,0 +1,95 @@
+"""CPU unit test for the optional direct-acceptance (LK) loss in DSpark.
+
+Runs without a GPU or a real model: it mocks a ``DSparkForwardOutput`` and checks
+  (1) accept_loss_alpha=0 is exactly the pre-existing CE(+L1+confidence) loss (zero regression);
+  (2) the accept term is differentiable (gradient reaches draft_logits);
+  (3) it rewards a draft that matches the target (lower loss when draft==target logits).
+
+Usage: python scripts/test_accept_loss.py
+"""
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+# metrics may accumulate to global rank state; make it a no-op for the unit test.
+import deepspec.utils.metrics as _metrics  # noqa: E402
+_metrics.add_metric = lambda *a, **k: None
+import deepspec.modeling.dspark.loss as L  # noqa: E402
+L.add_metric = lambda *a, **k: None
+from deepspec.modeling.dspark.common import DSparkForwardOutput  # noqa: E402
+
+
+def _init_dist():
+    if not dist.is_initialized():
+        os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+        os.environ.setdefault("MASTER_PORT", "29591")
+        dist.init_process_group("gloo", rank=0, world_size=1)
+
+
+def _mock(B=1, A=2, block=4, V=32, draft_eq_target=False, seed=0):
+    g = torch.Generator().manual_seed(seed)
+    target_logits = torch.randn(B, A, block, V, generator=g)
+    draft = target_logits.clone() if draft_eq_target else torch.randn(B, A, block, V, generator=g)
+    draft = draft.detach().requires_grad_(True)
+    target_ids = target_logits.argmax(dim=-1)
+    return draft, DSparkForwardOutput(
+        draft_logits=draft,
+        target_ids=target_ids,
+        eval_mask=torch.ones(B, A, block, dtype=torch.bool),
+        block_keep_mask=torch.ones(B, A, dtype=torch.bool),
+        confidence_pred=None,
+        aligned_target_logits=target_logits,
+    )
+
+
+def _loss(out, accept_alpha):
+    return L.compute_dspark_loss(
+        outputs=out, loss_decay_gamma=None,
+        ce_loss_alpha=1.0, l1_loss_alpha=0.0, confidence_head_alpha=0.0,
+        accept_loss_alpha=accept_alpha,
+    )
+
+
+def main():
+    _init_dist()
+
+    # (1) zero regression: accept_loss_alpha=0 == the CE-only loss to the bit
+    _, o0 = _mock(seed=1)
+    base = _loss(o0, 0.0)
+    _, o0b = _mock(seed=1)
+    off = _loss(o0b, 0.0)
+    assert torch.allclose(base, off), (base, off)
+    print(f"(1) zero-regression OK  (CE-only loss={base.item():.4f})")
+
+    # (2) differentiable: grad reaches draft_logits, and the accept term changes the loss
+    draft, o1 = _mock(seed=2)
+    withacc = _loss(o1, 0.5)
+    _, o1b = _mock(seed=2)
+    ceonly = _loss(o1b, 0.0)
+    assert not torch.allclose(withacc, ceonly), "accept term did not change the loss"
+    withacc.backward()
+    assert draft.grad is not None and torch.isfinite(draft.grad).all() and draft.grad.abs().sum() > 0
+    print(f"(2) differentiable OK   (loss +accept={withacc.item():.4f} vs CE-only={ceonly.item():.4f}; grad flows)")
+
+    # (3) rewards matching: a draft equal to target has a LOWER (more negative) accept term
+    dm, om = _mock(seed=3, draft_eq_target=True)   # draft == target
+    dr, orr = _mock(seed=3, draft_eq_target=False)  # random draft
+    # isolate the accept term (ce_alpha=0, accept_alpha=1)
+    acc_match = L.compute_dspark_loss(outputs=om, loss_decay_gamma=None,
+                                      ce_loss_alpha=0.0, l1_loss_alpha=0.0,
+                                      confidence_head_alpha=0.0, accept_loss_alpha=1.0)
+    acc_rand = L.compute_dspark_loss(outputs=orr, loss_decay_gamma=None,
+                                     ce_loss_alpha=0.0, l1_loss_alpha=0.0,
+                                     confidence_head_alpha=0.0, accept_loss_alpha=1.0)
+    assert acc_match.item() < acc_rand.item(), (acc_match.item(), acc_rand.item())
+    print(f"(3) rewards matching OK (accept-loss match={acc_match.item():.4f} < random={acc_rand.item():.4f})")
+
+    print("\nALL PASS")
+
+
+if __name__ == "__main__":
+    main()