Why fused KV injection instead of independent per-layer injection?

[shreyas269]

Hi @jianc99 , great work!!

Currently, DFlash extracts hidden states from 5 target layers, concatenates them (5 × d_target), projects down to d_target via a learned projection, and injects the same fused vector into every draft layer's KV cache.

Have you considered a simpler alternative? Instead of fusing, just map a subset (the same 5 layers) of target layer's hidden state directly to the corresponding draft layer? Eg. Let's say target has 40 layers and draft has 5 layers, then

Draft layer 1 ← Target layer 1 hidden state
Draft layer 2 ← Target layer 8 hidden state
...and so on.

Since the draft model shares d_target with the target model, no projection is needed at all...the hidden states can go directly through each draft layer's W_K, W_V. This avoids the lossy 5 × d_target → d_target compression, removes the extra projection layer entirely, and still doesn't restrict information access across target layers since draft layers are stacked and information from earlier injections propagates through the residual stream.

Curious if this was ablated and if there were specific reasons it underperformed...

[jianc99]

@shreyas269 Thanks for the thoughtful suggestion! This is a very reasonable alternative, and we did ablate it.

Our main motivation for fusing hidden states first is to ensure that every draft layer has access to information from multiple target depths. Shallow and deep target layers encode different signals, and exposing each draft layer to this mixed information turns out to be important for draft quality.

We compared this against a direct layer-aligned mapping (feeding different target layer hidden states into corresponding draft layers without fusion). These two models are trained on our 100K ablation dataset. As shown in the plots, the fused variant (blue) consistently achieves higher per-position accuracy and higher acceptance length than the non-fused variant (purple). Empirically, restricting each draft layer to a single target layer underperforms, and residual propagation alone does not fully compensate for the missing cross-depth context.

Additionally, we believe that fusing target hidden states helps decouple the number of target hidden states used from the number of draft layers, which provides more flexibility and a larger design space for future draft architectures. Based on these results, we believe fusing target hidden states before injection is the better design choice, despite the extra projection.

[shreyas269]

Thanks for the response!

I agree that decoupling the number of target hidden states from the number of draft layers provides more flexibility. The ablations are interesting...clearly the fused KV injection seems to have higher acceptance length. I asked my original question with the intuition that we kind of want the draft model to mimic the target model as a lightweight replacement, so layer-aligned mapping for KV from target to draft model seemed a natural choice here.

Also you mention about the acceptance length scaling up as we increase the number of target hidden features. This also comes at the cost of additional training time (which you did mention in the paper) as well as additional I/O during inference. I believe the number of projection params will just be #hidden_layers (which is 5) since it's a linear projection.

These ablations are insightful! Thanks for sharing them! 💯

[catnanami]

@jianc99 Hi！I am curious about whether you have ablated this with different numbers of hidden layer？

@shreyas269 Thanks for the thoughtful suggestion! This is a very reasonable alternative, and we did ablate it.

Our main motivation for fusing hidden states first is to ensure that every draft layer has access to information from multiple target depths. Shallow and deep target layers encode different signals, and exposing each draft layer to this mixed information turns out to be important for draft quality.

We compared this against a direct layer-aligned mapping (feeding different target layer hidden states into corresponding draft layers without fusion). These two models are trained on our 100K ablation dataset. As shown in the plots, the fused variant (blue) consistently achieves higher per-position accuracy and higher acceptance length than the non-fused variant (purple). Empirically, restricting each draft layer to a single target layer underperforms, and residual propagation alone does not fully compensate for the missing cross-depth context.

Additionally, we believe that fusing target hidden states helps decouple the number of target hidden states used from the number of draft layers, which provides more flexibility and a larger design space for future draft architectures. Based on these results, we believe fusing target hidden states before injection is the better design choice, despite the extra projection.