[DFlash/DSpark] Optional direct-acceptance (LK) loss#58
Open
zhangwenhe1007 wants to merge 1 commit into
Open
Conversation
Adds a differentiable direct-acceptance objective that trains the drafter to
maximise expected accepted length E[tau]=sum_k prod_{j<=k} alpha_j directly,
instead of only the per-token cross-entropy proxy. It makes the acceptance
quantity already logged as the tau_probabilistic metric trainable.
- deepspec/modeling/dspark/loss.py: accept_loss term (gated on accept_loss_alpha>0)
- deepspec/trainer/dspark_trainer.py: wire accept_loss_alpha (getattr default 0.0)
- config/dflash/dflash_qwen3_8b.py: documented accept_loss_alpha knob (default 0.0)
- scripts/test_accept_loss.py: CPU unit test (zero-regression, differentiable, rewards matching)
Implementation prepared with AI assistance.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Current drafters train on per-token cross-entropy, but the quantity that actually determines the speedup is the acceptance length. The repo already computes the expected acceptance length (E[tau] = 1 + sum_k prod_{j<=k} alpha_j) and logs it as the
tau_probabilisticmetric, but it sits undertorch.no_grad()so it's never optimized. This PR makes it an optional loss, such that you can train the drafter to maximize acceptance directly instead of only through the CE proxy.It's gated behind
accept_loss_alpha(default 0.0), so existing configs are untouched. The term is just the negative expected accepted length, differentiable through the per-position acceptance rate, and it needsaligned_target_logits(same as the L1 and confidence terms).Files:
The test runs on CPU with no model and verifies three things: accept_loss_alpha=0 gives the exact same loss as before, the term is differentiable (gradient reaches the draft logits), and it rewards a draft that matches the target (the loss goes to -4.0 when all four block positions accept with probability 1).
In my own DFlash runs, a direct acceptance loss like the one in this PR improved position-1 acceptance by about 6% over CE at matched data.