论文：OPT-Tree: Speculative Decoding with Adaptive Draft Tree Structure
链接：https://arxiv.org/abs/2406.17276
代码：https://github.com/Jikai0Wang/OPT-Tree

1.背景

目前主流的大模型大多是自回归模型(Autoregressive models)，其“一步一词”的生成模式带来了较大的推理开销，制约了其在各种场景中的应用。

投机解码通过”小模型起草(Drafting)，大模型验证(Verification)“的方式能够在大模型的一步推理中生成多个有效tokens，实现了无损的模型解码加速。

在投机解码中，小模型的性能决定了其生成的draft的质量，起草的质量越高，平均验证通过的长度也就越大。另一方面，draft的数据结构也是影响算法性能的一个重要因素。

2.方法

以往的工作大多采用序列或是启发式的固定树结构的draft。序列形式的draft存在较多的前缀冗余，即多个draft序列在同一个位置上可能出现多个相同的token。固定的树结构虽然避免了这种冗余，但依然不是有限预算下最优的draft结构。

考虑到理论上最优的树结构因当在解码的每一步都可能是不同的，本文提出了一种自适应且可扩展的draft结构——OPT-Tree。在给定结点个数的情况下，OPT-Tree能够找到每一步中最大化近似验证通过长度期望E(A)的一个draft树结构。E(A)的定义如下：

3.实验结果

在不同目标模型与draft模型的组别中，OPT-Tree的无论是平均验证通过长度(MAL)还是吞吐速度均优于现有的draft结构。

文中还展示了一个以7B LLAMA作为draft模型加速70B LLAMA的一个实例：

其中蓝色的文本均由draft模型生成，再由大模型进行并行验证并通过，红色的文本是模型验证的副产物(实际由大模型生成)。这一例子中的平均验证通过长度为9.34。

4.总结

OPT-Tree提供了一种自适应的draft树结构，适用于各种自回归的draft模型，相信其优越的性能也将为后续的解码加速相关工作带来启发。