search

SafeAILab / EAGLE

Official Implementation of EAGLE
https://arxiv.org/pdf/2406.16858
Apache License 2.0
622 stars 59 forks source link

Question on draft process #58

Closed cyLi-Tiger closed 3 months ago

cyLi-Tiger commented 3 months ago

I had a question on the sampling strategy during building the token tree. I noticed here you sampled the next layer using torch.multinomial, why not use torch.topk instead? Cause it seems that choosing the top k likely selected tokens are more reasonable. And SpecInfer also use the top k candidates to expand a token tree.

Liyuhui-12 commented 3 months ago

As long as the sampling probabilities are recorded, both sampling methods are correct. This is because speculative sampling does not have any requirements for the draft distribution q. Top-k sampling is equivalent to repeating argmax multiple times. When the draft distribution q and the target distribution p are highly consistent, sampling is better than argmax. For example, if p=(0.6,0.4) and q=(0.6,0.4), the acceptance rate of sampling from q is 1. When using argmax, it is equivalent to sampling from the distribution (1.0,0.0), and the acceptance rate is 0.6.