search

Infini-AI-Lab / Sequoia

scalable and robust tree-based speculative decoding algorithm
282 stars 29 forks source link

Question on tree search algorithm #15

Closed cyLi-Tiger closed 3 months ago

cyLi-Tiger commented 3 months ago

Here, the max_branch equals to K + 1, K refers to Algorithm 5 in your paper. The K + 1 dimension of p represents the percentage of accept none from the draft model according to your code in test_accept.py. So there is gap between your code and Algorithm 5? Correct me if I misunderstood anything!

dreaming-panda commented 3 months ago

In tree_search.py, vector p is indiced from 1. For example, p = [0.0, 0.4803, 0.1104, 0.0576, 0.0373, 0.0265, 0.0211, ....]. The first element is 0 and acts as a padding.

cyLi-Tiger commented 3 months ago

In tree_search.py, vector p is indiced from 1. For example, p = [0.0, 0.4803, 0.1104, 0.0576, 0.0373, 0.0265, 0.0211, ....]. The first element is 0 and acts as a padding.

Indeed, p[0] == 0. But len(p) == K + 2 and max_branch == K + 1, I'm asking about the last dimension of p, which

represents the percentage of accept none from the draft model according to your code in test_accept.py.

Since in Algorithm 5 you consider max_branch range from 0 to K, and in the code is 0 to K + 1. For example, if K=32, the output_branch_prob in test_accept.py has 34 dimensions and p.shape[0] == 34

dreaming-panda commented 3 months ago

I think you are correct, and I just fixed it by p = p [:-1] .

Previously, I do not notice that because usually we will not use a tree as wide as 32. Thank you!