Add D-PACE loss implementation for D-Flash training

scripts/train.py

@@ -961,6 +961,20 @@ def parse_args():
         default=4.0,
         help="Decay gamma for DFlash/DSpark loss weighting (default: 4.0)",
     )
+    # D-Pace specific arguments (loss weight option + smoothing)
+    parser.add_argument(
+        "--per-position-loss-weight",
+        choices=["fixed-exp-decay", "dpace"],
+        default="fixed-exp-decay",
+        help="Per-position loss weight option for D-PACE support"
+        "default: fixed-exp-decay",
+    )
+    parser.add_argument(
+        "--dpace-alpha",
+        type=float,
+        default=0.5,
+        help="Smoothing constant for D-PACE loss (default: 0.5)",
+    )
     # DSpark-specific arguments (sequential Markov head + confidence head).
     parser.add_argument(
         "--markov-rank",
@@ -1129,6 +1143,13 @@ def parse_args():
     provided = explicitly_provided_dests(parser, DECODER_SHAPING_FLAGS)
     validate_draft_init_args(parser, args, provided)
     resolve_loss_config(args.loss_fn)
+
+    if args.per_position_loss_weight == "dpace":
+        if args.loss_fn != "ce":
+            parser.error("--per-position-loss-weight=dpace requires --loss-fn=ce")
+        if not 0.0 < args.dpace_alpha <= 1.0:
+            raise ValueError(f"alpha must be in (0, 1], got {args.dpace_alpha}")
+
     return args
 
 

src/speculators/models/dflash/core.py

@@ -207,10 +207,16 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:
         loss_config = resolve_loss_config(kwargs["loss_fn"])
         gamma = kwargs.get("dflash_decay_gamma", 4.0)
         max_anchors = kwargs.get("max_anchors", 3072)
+        per_position_loss_weight = kwargs.get(
+            "per_position_loss_weight", "fixed-exp-decay"
+        )
+        dpace_alpha = kwargs.get("dpace_alpha", 0.5)
         shared = {
             "loss_config": loss_config,
             "gamma": gamma,
             "max_anchors": max_anchors,
+            "per_position_loss_weight": per_position_loss_weight,
+            "dpace_alpha": dpace_alpha,
         }
         return dict(shared), dict(shared)
 
@@ -393,6 +399,8 @@ def forward(
         loss_config: LossConfig | None = None,
         gamma: float = 4.0,
         max_anchors: int = 3072,
+        per_position_loss_weight: str = "fixed-exp-decay",
+        dpace_alpha: float = 0.5,
         **kwargs,
     ):
         _, logits, targets, aligned_loss_mask, _ = self._backbone_forward(
@@ -412,6 +420,8 @@ def forward(
             self.block_size,
             gamma=gamma,
             loss_config=loss_config,
+            per_position_loss_weight=per_position_loss_weight,
+            dpace_alpha=dpace_alpha,
         )
         draft_tokens = torch.argmax(logits, dim=-1)
 

src/speculators/models/dflash/metrics.py

@@ -10,6 +10,7 @@
     compound_loss,
     compute_accuracy_multi_step,
     dflash_loss_decay,
+    dpace_loss_decay,
     kl_div_loss,
 )
 
@@ -23,6 +24,8 @@ def compute_metrics(
     block_size: int = 1,
     gamma: float = 4.0,
     loss_config: LossConfig | None = None,
+    per_position_loss_weight: str = "fixed-exp-decay",
+    dpace_alpha: float = 0.5,
 ) -> tuple[torch.Tensor, dict]:
     """Compute loss and accuracy metrics for draft model predictions.
 
@@ -33,6 +36,8 @@ def compute_metrics(
         block_size: Block size for per-position metrics
         gamma: Temperature for exponential decay in loss weighting
         loss_config: Mapping of ``{name: (loss_fn, weight)}``
+        per_position_loss_weight: Weighting option for per-position block-drafting loss
+        dpace_alpha: Smoothing constant for D-Pace loss weighting
 
     Returns:
         Tuple of (loss, metrics_dict) where metrics_dict contains:
@@ -46,13 +51,23 @@ def compute_metrics(
     pos_idx = torch.arange(seq_len, device=logits.device) % block_size
     pos_idx = pos_idx.unsqueeze(0)  # shape: [1, T]
 
+    if per_position_loss_weight == "dpace":
+        decay_fn = partial(
+            dpace_loss_decay,
+            loss_mask=loss_mask,
+            block_size=block_size,
+            dpace_alpha=dpace_alpha,
+        )
+    else:
+        decay_fn = partial(dflash_loss_decay, gamma=gamma)
+
     loss, term_losses = compound_loss(
         logits,
         targets,
         loss_mask,
         pos_idx,
         loss_config=loss_config,
-        decay_fn=partial(dflash_loss_decay, gamma=gamma),
+        decay_fn=decay_fn,
     )
 
     pred_ids = torch.argmax(logits, dim=-1)

src/speculators/models/dspark/core.py

@@ -84,11 +84,17 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:
         gamma = kwargs.get("dflash_decay_gamma", 4.0)
         max_anchors = kwargs.get("max_anchors", 3072)
         confidence_head_alpha = kwargs.get("confidence_head_alpha", 1.0)
+        per_position_loss_weight = kwargs.get(
+            "per_position_loss_weight", "fixed-exp-decay"
+        )
+        dpace_alpha = kwargs.get("dpace_alpha", 0.5)
         shared = {
             "loss_config": loss_config,
             "gamma": gamma,
             "max_anchors": max_anchors,
             "confidence_head_alpha": confidence_head_alpha,
+            "per_position_loss_weight": per_position_loss_weight,
+            "dpace_alpha": dpace_alpha,
         }
         return dict(shared), dict(shared)
 
@@ -105,6 +111,8 @@ def forward(
         gamma: float = 4.0,
         max_anchors: int = 3072,
         confidence_head_alpha: float = 1.0,
+        per_position_loss_weight: str = "fixed-exp-decay",
+        dpace_alpha: float = 0.5,
         **kwargs,
     ):
         hidden, logits, targets, aligned_loss_mask, anchored_block_indices = (
@@ -167,6 +175,8 @@ def forward(
             loss_config=loss_config or _DEFAULT_LOSS_CONFIG,
             gamma=gamma,
             confidence_head_alpha=confidence_head_alpha,
+            per_position_loss_weight=per_position_loss_weight,
+            dpace_alpha=dpace_alpha,
         )
         draft_tokens = torch.argmax(logits, dim=-1)
         return draft_tokens, loss, metrics

src/speculators/models/dspark/metrics.py

@@ -18,6 +18,7 @@
     compound_loss,
     compute_accuracy_multi_step,
     dflash_loss_decay,
+    dpace_loss_decay,
 )
 
 __all__ = [
@@ -31,13 +32,15 @@ def _masked_decayed_mean(
     elementwise: torch.Tensor,  # [1, T]
     loss_mask: torch.Tensor,  # [1, T]
     pos_idx: torch.Tensor,  # [1, T]
-    decay_fn: Callable[[torch.Tensor], torch.Tensor] | None,
+    decay_fn: Callable[..., torch.Tensor] | None,
 ) -> torch.Tensor:
     """Masked, optionally position-decayed mean of a precomputed per-position term."""
     loss_mask = loss_mask.to(elementwise.dtype)
     weighted = elementwise * loss_mask
     if decay_fn is not None:
-        weighted = weighted * decay_fn(pos_idx.to(weighted.dtype))
+        weighted = weighted * decay_fn(
+            pos_idx.to(weighted.dtype), elementwise_loss=elementwise
+        )
     denominator = loss_mask.sum(dim=1) + _EPS
     return (weighted.sum(dim=1) / denominator).mean()
 
@@ -51,13 +54,23 @@ def compute_metrics(
     loss_config: LossConfig,
     gamma: float = 4.0,
     confidence_head_alpha: float = 1.0,
+    per_position_loss_weight: str = "fixed-exp-decay",
+    dpace_alpha: float = 0.5,
 ) -> tuple[torch.Tensor, dict]:
     """Compute the DSpark loss and a metrics dict (``*_sum``/``*_total`` pairs)."""
 
     device = logits.device
     seq_len = logits.shape[1]
     pos_idx = (torch.arange(seq_len, device=device) % block_size).unsqueeze(0)
-    decay_fn = partial(dflash_loss_decay, gamma=gamma)
+    if per_position_loss_weight == "dpace":
+        decay_fn = partial(
+            dpace_loss_decay,
+            loss_mask=loss_mask,
+            block_size=block_size,
+            dpace_alpha=dpace_alpha,
+        )
+    else:
+        decay_fn = partial(dflash_loss_decay, gamma=gamma)
 
     loss, term_losses = compound_loss(
         logits, targets, loss_mask, pos_idx, loss_config=loss_config, decay_fn=decay_fn

src/speculators/models/metrics.py

@@ -244,7 +244,7 @@ def lk_hybrid_loss(
     return elementwise_loss  # noqa: RET504
 
 
-def dflash_loss_decay(pos_idx: torch.Tensor, gamma: float):
+def dflash_loss_decay(pos_idx: torch.Tensor, gamma: float, **_kwargs):
     """Compute DFlash-style exponential decay weights per position.
 
     Position 0 gets weight 0, position 1 gets weight 1, and subsequent positions
@@ -265,7 +265,7 @@ def dflash_loss_decay(pos_idx: torch.Tensor, gamma: float):
     return decay_mult  # noqa: RET504
 
 
-def exp_loss_decay(pos_idx: torch.Tensor, gamma: float):
+def exp_loss_decay(pos_idx: torch.Tensor, gamma: float, **_kwargs):
     """Compute simple exponential decay weights as gamma^pos_idx.
 
     Args:
@@ -278,6 +278,57 @@ def exp_loss_decay(pos_idx: torch.Tensor, gamma: float):
     return gamma**pos_idx
 
 
+def dpace_loss_decay(
+    pos_idx: torch.Tensor,  # noqa: ARG001
+    loss_mask: torch.Tensor,
+    block_size: int,
+    dpace_alpha: float,
+    elementwise_loss: torch.Tensor,
+    **_kwargs,
+):
+    """
+    Per-position block-drafting loss weight based on D-PACE
+
+    Args:
+        elementwise_loss: requires to be cross-entropy loss, negative log-likelihood
+            of per-position confidence
+        dpace_alpha: confidence smoothing constant
+
+    Returns:
+        Decay multiplier tensor with same shape as pos_idx.
+    """
+    with torch.no_grad():
+        # convert CE to per-position confidence
+        q = torch.exp(-elementwise_loss).float()
+
+        # reshape loss to [num_anchors, block_size]
+        # for intra-block cumulative multiplication
+        if q.shape[1] % block_size != 0:
+            raise ValueError(
+                f"q.shape[1] ({q.shape[1]}) must be divisible by "
+                f"block_size ({block_size})"
+            )
+        num_anchors = q.shape[1] // block_size
+        q = q.reshape(num_anchors, block_size)
+        mask = loss_mask.reshape(num_anchors, block_size).to(q.dtype)
+
+        # smoothed confidence for numerical stability
+        smooth = (1.0 - dpace_alpha) * q + dpace_alpha
+        smooth = torch.where(mask > 0, smooth, torch.ones_like(smooth))
+
+        # prefix cumulative production
+        prefix = torch.cumprod(smooth, dim=-1)
+
+        # suffix summation: flip -> cumsum -> flip
+        weight = torch.flip(
+            torch.cumsum(torch.flip(prefix * mask, dims=[-1]), dim=-1), dims=[-1]
+        )
+        weight = weight * mask
+
+    # reshape weight
+    return weight.reshape(1, -1)
+
+
 _LOSS_FN_MAP: dict[str, Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = {
     "kl_div": kl_div_loss,
     "rkl": reverse_kl_div_loss,
@@ -360,7 +411,7 @@ def compound_loss(
     loss_mask: torch.Tensor,
     pos_idx: torch.Tensor,
     loss_config: LossConfig,
-    decay_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+    decay_fn: Callable[..., torch.Tensor] | None = None,
 ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
     """Compute a weighted sum of loss terms.
 
@@ -377,7 +428,12 @@ def compound_loss(
     multi = len(loss_config) > 1
     for name, (fn, weight) in loss_config.items():
         term = loss_function(
-            logits, targets, loss_mask, pos_idx, loss_fn=fn, decay_fn=decay_fn
+            logits,
+            targets,
+            loss_mask,
+            pos_idx,
+            loss_fn=fn,
+            decay_fn=decay_fn,
         )
         if multi:
             term_losses[f"{name}_loss"] = term.detach()
@@ -391,7 +447,7 @@ def loss_function(
     loss_mask: torch.Tensor,  # shape: [1, seq_len]
     pos_idx: torch.Tensor,  # shape: [1, seq_len]
     loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = kl_div_loss,
-    decay_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+    decay_fn: Callable[..., torch.Tensor] | None = None,
 ):
     """Compute masked, optionally position-decayed training loss.
 
@@ -412,7 +468,9 @@ def loss_function(
     elementwise_loss = elementwise_loss * loss_mask
 
     if decay_fn is not None:
-        decay_mult = decay_fn(pos_idx.to(elementwise_loss.dtype))
+        decay_mult = decay_fn(
+            pos_idx.to(elementwise_loss.dtype), elementwise_loss=elementwise_loss
+        )
         elementwise_loss = elementwise_loss * decay_mult
 
     denominator = loss_mask.sum(dim=1) + _EPS