DynaDFlash: Dynamic Feature Projection for Block Diffusion Speculators

Introduction

Speculative decoding uses a smaller draft model to predict tokens that are then verified by a larger target model. DFlash does this by training a block diffusion model to draft multiple tokens in parallel, using hidden states from the target model as additional context that is injected at each layer.

DFlare builds on this idea by learning different combinations of target hidden states for different draft layers. However, these combinations remain fixed for each layer after training. We wanted to test whether allowing the weights to change depending on the input context could further improve drafting.

We modify DFlash by adding a small gate that computes weights for the target hidden states before they are projected. We then train both the original DFlash architecture and our modified version under the same conditions and compare their acceptance lengths and decoding throughput.

Background

Autoregressive LLMs generate tokens one at a time. Each token requires a new forward pass through the target model, which slows down generation. Speculative decoding allows us to produce more tokens for each forward pass of the target model. A smaller model first drafts several tokens, and the target model checks them together in a single forward pass.

Under greedy decoding, a draft token is accepted if it matches the highest likelihood token the target would have produced. Tokens are accepted until the first mismatch, at which point the target produces the next token instead. The final output is therefore the same as ordinary greedy decoding. For sampling, speculative decoding uses a different acceptance formula to preserve the target model's output distribution.

DFlash reduces drafting time by predicting the next 15 tokens in parallel from a block of masked positions, rather than generating each draft token one at a time. The target then verifies the block. This makes both drafting time and acceptance length important when comparing draft models.

Models

DFlash takes hidden states from layers 2, 10, 18, 26, and 34 of the target model, concatenates them, and projects them into the draft model.

The gate takes the hidden state from layer 34 and computes five weights, one for each of the selected target layers. Each hidden state is multiplied by its corresponding weight before concatenation and projection.

Unlike the fixed projection in DFlash, these weights are computed separately for each token, allowing the contribution of each target layer to depend on the context.

The gate introduces only 20,480 additional parameters.

Evaluation

Training.

We train both models using a frozen Qwen3-8B target and the same 792,000 training examples in the same order. Both use DFlash's original training methodology with a block size of 16.

Training runs for one epoch, consisting of 24,750 updates with a global batch size of 32 on two H100 GPUs. We use 990 warmup steps, a peak learning rate of 6×10⁻⁴, and cosine learning rate decay. We evaluate the final checkpoints after checking their training histories.

Math evaluation.

After observing stronger performance on math prompts, we conduct a separate test using 700 held-out prompts from MetaMathQA and Orca-Math. Of these, 647 come from MetaMathQA and 53 from Orca-Math.

For each prompt, we provide the question without its answer. The target model generates the first token, and the draft model predicts the next 15 tokens. We then measure how many tokens match the target model's greedy predictions, including the first token generated by the target. We count acceptances from the beginning of the draft and stop at the first mismatch.

Serving.

We also measure decoding throughput using SGLang 0.5.18 on one A100 GPU. We select 24 fixed MetaMathQA prompts and generate 64 tokens per request using greedy decoding.

We test both models with eight concurrent requests. Each experiment is repeated three times using new servers, alternating the order in which the models are tested. Throughput is measured in output tokens per second.

Results

On the 700 math prompts, the dynamic model achieves a mean accepted length of 3.13 tokens, compared to 2.98 for the original DFlash model. This is an improvement of about 5%, or 0.15 additional tokens per prompt.

The 95% confidence interval for the difference is 0.07 to 0.22 tokens, with p<0.001.

Most of this improvement comes from the 647 MetaMathQA prompts, where the dynamic model accepts an additional 0.16 tokens on average. We do not observe an improvement on the smaller set of 53 Orca-Math prompts.

The results are different on our full evaluation, where the original DFlash model and the dynamic model have no statistically significant difference in acceptance length. This suggests that the benefits of dynamic weighting may depend on the type of input rather than improving drafting across all tasks.

We also observe a throughput improvement when serving eight requests in parallel. The dynamic model generates an average of 1004.5 output tokens per second compared to 965.0 for DFlash, an improvement of 4.1%.

Discussion

Our results show that dynamic context projection can improve DFlash's acceptance length on math prompts, but the improvement does not extend to all types of prefixes. In particular, the dynamic model performs better on MetaMathQA while not significantly differing in performance from the baseline for the full evaluation.

We also looked at the weights produced by the gate to understand whether certain target layers were more useful for math. On the 700 math prompts, layers 2 had the largest mean gate, and layer 26 had the second largest. However, due to differences in the magnitude of the hidden features for these layers, layer 26 contributed the most to the projection, followed layer 34.

A key limitation is that both models were trained independently and do not match the acceptance length of the published DFlash model by z-lab. This makes the observed improvement in acceptance length less reliable as both models fall below the acceptance length of other DFlash drafters.

Conclusion

We tested whether dynamically weighting target hidden states before projection could improve DFlash's draft model. Our modified model achieves a 5% higher mean accepted length on 700 math prompts and a 4.1% throughput improvement in our serving test.

However, these improvements are not consistent across datasets, and we could not establish a clear relationship between dynamic context to the drafter and model performance. Further experimentation may be required with different methods to inject dynamic information from the target model to the drafter to demonstrate a gain across various contexts.