投机采样论文整理

akumaexs

1. Speculative Decoding

绿色表示the approximation model(小模型)的输出被target model(大模型)所接受，红色表示reject，蓝色表示correction

用小模型$M_q$做自回归采样连续生成 $γ$ 个tokens。
把生成的γ个tokens和前缀拼接一起送进大模型$M_p$执行一次forwards。
使用大、小模型logits结果做比对，如果发现某个token小模型生成的不好，重新采样这个token。重复步骤1。
如果小模型生成结果都满意，则用大模型采样下一个token。重复步骤1。
第2步，将γ个tokens和前缀拼成一起作为大模型输入，和自回归相比，尽管计算量一样，但是γ个tokens可以同时参与计算，计算访存比显著提升。
第3步，如何评价一个token生成的不好？如果q(x) > p(x)（p，q表示在大小模型采样概率，也就是logits归一化后的概率分布）则以$1-p(x)/q(x)$为概率拒绝这个token的生成，从一个新的概率分布$$p’(x) = norm(max(0, p(x) − q(x)))$$中重新采样一个token。

akumaexs

2. Medusa

Speculative Decoding 难以获取draft模型，Medusa通过复用原始LLM的权重，训练参数规模较小的增量权重用于预测 Draft Token，这种在生成 Draft Token 时复用原始 LLM 权重的方式也叫 Self Drafting

有人认为Medusa的方法并不符合投机采样的思想

LM_Head是原始模型的权重，用于预测Next token，Medusa Head是增量权重，负责在Draft phase预测Draft Tokens。在多个 Medusa Head 中，Medusa Head 1 负责预测 Next Next Token，Medusa Head 2 负责预测 Next Next Next Token

验证Draft

Draft的验证路径其实是一个有向无环图，因此验证开销很大。
Medusa提出Tree attention，将所有路径展开为一维的序列，将这个序列作为一个序列输入放到模型中运行。其中 Attention 阶段 Token 之间会进行注意力计算，为了避免不同路径之间的 Token 计算注意力，需要将其他 Token 加上 Mask，不让某一路径的 Token “看到” 其他路径的 Token。
此外，为了避免不必要的overhead，通过先验的方式对 Draft 树进行剪枝，将接收概率低的路径提前剪掉，缩短验证序列长度。

存在的问题

Medusa 采用了隐状态并行预测多步 Draft Tokens，在预测间隔 Tokens 时（即 Next Token 以后的 Draft Token），Draft Token 并不知道上一个 Token，进而影响了准确率

akumaexs

3. EAGLE

Speculative Decoding是Token-level 采样，Medusa是Feature-level采样(last hidden state)。
为了提高 Token 生成的稳定性，在生成当前 Token的隐藏层特征时，可以融合当前 Token 的 Embedding 以及上一个 Token 的隐藏层特征，保证 Draft Token生成时具备足够的信息，增强 Token 生成效果。

Logits: $p_{j+1}=LM\_Head(f_j)$

token生成的顺序为：$T{1:j} \rightarrow E{1:j} \rightarrow f_j \rightarrow p{j+1} \rightarrow t{j+1}$
Draft 模型是一个同构的轻量 LLM，其中 Embedding 层和 LM Head 层均复用原始 LLM 模型，中间的 One Auto-regression Head，由一层 FC 层以及一层 Transformer Layer 组成。
AR Head 是唯一需要微调的网络层，训练成本也是极低。EAGLE 使用 Draft 模型通过自回归采样方式生成 Draft Tokens，并且为了提升接收率，在运行 AR Head 前会融合上一个 Token 的隐状态与当前 Token 的 Embedding层，并通过 AR Head 的 FC 层融合

在Draft阶段，自回归采样的第一轮输入是原始 LLM 模型生成的 Token，通过 TopK 采样生成第一轮的 Top $K_1$个Draft Token。该$K_1$个 Draft Token作为第二轮输入，对应生成 $K_1$组 Draft Tokens，每组通过采样生成 $K_2$ 个 Draft Token。下一轮会从当前轮中选择若干个 Draft Token 的输出作为输入，继续自回归采样，直至到达预设的运行轮次。