Introduction
- 作者提出 Context-Aware Dynamic Draft Tree,对 EAGLE 中固定的 draft tree 做出了改进,通过动态扩展 draft tree 有效节省 draft model 的推理时间,还能通过 token re-rank 给 LLM 提供更少更高质量的 draft tokens,从而在保证接收率的前提下降低 LLM 的验证时间,相比 EAGLE 有了明显的加速
Method
- Expansion Phase. 在扩展 draft tree 时,只选择最深层上累积接收概率最高的 top-
k
k
k tokens 进行扩展,从而降低 draft model 的推理开销,其中累积接收概率为根节点到当前节点的接收概率之积,接收概率可以用 token 置信度近似
V
i
=
∏
t
j
∈
Path
(
root
,
t
i
)
p
j
≈
∏
t
j
∈
Path
(
root
,
t
i
)
c
j
V_i=\prod_{t_j\in\text{Path}(\text{root},t_i)}p_j\approx\prod_{t_j\in\text{Path}(\text{root},t_i)}c_j
Vi=∏tj∈Path(root,ti)pj≈∏tj∈Path(root,ti)cj (作者发现 draft model is well-calibrated);例如下图中,只选择 current layer (orange blocks) 进行扩展
- Reranking Phase. draft tree 扩展完成后,作者从 token tree 中选出累积接收概率的 top-
m
m
m tokens 交给 LLM 验证,如果两个 token 累计接收概率相同,则优先选择位于更浅层的 token;这样选出的 top-
m
m
m tokens 仍然可以形成一棵树,可以用 tree attention 进行验证
Experiments
- Effectiveness.
- Ablation Study. 这里的 w/o value 指扩展 draft tree 时不使用累积接收概率而是直接使用单个 token 的 confidence