def compute_metrics(
logits: torch.Tensor, # [1, T, draft_vocab_size] (Markov-corrected)
targets: torch.Tensor, # [1, T, draft_vocab_size]
confidence_logits: torch.Tensor | None, # [1, T] or None
loss_mask: torch.Tensor, # [1, T]
block_size: int,
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)
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
)
# Analytical per-position acceptance rate = distributional overlap.
with torch.no_grad():
draft_p = softmax(logits.float(), dim=-1)
target_p = softmax(targets.float(), dim=-1)
accept_rate = torch.minimum(draft_p, target_p).sum(dim=-1) # [1, T]
# Per-block cumulative acceptance product over the draft slots (slot 0
# is the anchor), shared by the accept-length and calibration metrics.
num_blocks = seq_len // block_size
accept_blocks = accept_rate.view(num_blocks, block_size)
draft_mask = loss_mask.to(accept_rate.dtype).view(num_blocks, block_size)[:, 1:]
accept_prefix = (accept_blocks[:, 1:] * draft_mask).cumprod(dim=-1)
metrics: dict[str, Any] = {}
if confidence_logits is not None:
c_star = accept_rate.detach().to(confidence_logits.dtype)
bce = binary_cross_entropy_with_logits(
confidence_logits, c_star, reduction="none"
) # [1, T]
conf_loss = _masked_decayed_mean(bce, loss_mask, pos_idx, decay_fn)
loss = loss + confidence_head_alpha * conf_loss
with torch.no_grad():
mask_f = loss_mask.to(accept_rate.dtype)
mask_total = mask_f.sum().clamp_min(1.0)
conf_prob = confidence_logits.float().sigmoid()
metrics["confidence_loss_sum"] = conf_loss.detach().clone()
metrics["confidence_loss_total"] = torch.ones((), device=device)
metrics["confidence_abs_error_sum"] = (
(conf_prob - accept_rate).abs() * mask_f
).sum()
metrics["confidence_abs_error_total"] = mask_total
# Mean predicted vs. observed acceptance — a calibration sanity check.
metrics["confidence_pred_mean_sum"] = (conf_prob * mask_f).sum()
metrics["confidence_pred_mean_total"] = mask_total
# Calibration of the cumulative acceptance product, which is what
# dynamic draft-length thresholding consumes (signed pred - target).
conf_prefix = (
conf_prob.view(num_blocks, block_size)[:, 1:] * draft_mask
).cumprod(dim=-1)
metrics["confidence_cumprod_bias_sum"] = (
(conf_prefix - accept_prefix) * draft_mask
).sum()
metrics["confidence_cumprod_bias_total"] = draft_mask.sum().clamp_min(1.0)
ones = torch.ones((), device=device)
metrics["loss_sum"] = loss.detach().clone()
metrics["loss_total"] = ones
for term_name, term_val in term_losses.items():
metrics[f"{term_name}_sum"] = term_val
metrics[f"{term_name}_total"] = ones
# Mean acceptance rate of the (Markov-corrected) drafter.
with torch.no_grad():
mask_f = loss_mask.to(accept_rate.dtype)
metrics["accept_rate_sum"] = (accept_rate * mask_f).sum()
metrics["accept_rate_total"] = mask_f.sum().clamp_min(1.0)
# Expected accepted draft length per block (DSpark's tau): the cumulative
# acceptance product summed over draft slots, plus the always-emitted anchor.
with torch.no_grad():
per_block_len = accept_prefix.sum(dim=-1) + 1.0
block_valid = (draft_mask.sum(dim=-1) > 0).to(accept_rate.dtype)
metrics["accept_len_sum"] = (per_block_len * block_valid).sum()
metrics["accept_len_total"] = block_valid.sum().clamp_min(1.0)
# Per-position greedy accuracy (position 0 is the anchor — excluded).
pred_ids = torch.argmax(logits, dim=-1)
target_ids = torch.argmax(targets, dim=-1)
correct_per_pos, total_per_pos = compute_accuracy_multi_step(
pred_ids, target_ids, loss_mask, pos_idx, block_size
)
metrics["full_acc_sum"] = correct_per_pos[1:].sum()
metrics["full_acc_total"] = total_per_pos[1:].sum()
for pos in range(1, block_size):
metrics[f"position_{pos}_acc_sum"] = correct_per_pos[pos]
metrics[f"position_{pos}_acc_total"] = total_per_pos[pos]
return loss, metrics