Extensive Reading
Author Info
Background
- Existing works only consider a token sequence generated by a single SSM for speculation which cannot align well with an LLM due to the model capacity gap between them.
- The probability of a successful alignment between the LLM and the speculated token sequence decays exponentially with the expected alignment length.
Challenges
- How to generate a token tree in a extremely large search space?
- How to verify the whole token tree in a single verfication pass?
Insights
Simultaneously consider a diversity of speculation candidates (instead of just one as in existing approaches) to maximize speculative performance.
These candidates are organized as a token tree.
Approaches
Token Tree Construction
- Expansion-based token tree construction
- Derive multiple tokens from an SSM within a single decoding step.
- 每次 decoding iteration 中选 top-k 就行,也就是在 token tree 中,对于一个父节点,选多个子节点出来。
- SpecInfer employs a static expansion strategy, [[EdgeLLM Fast On-Device LLM Inference With Speculative Decoding|EdgeLLM]] improves this by adopting a dynamic strategy - the branch with higher confidence is intended to be longer and stops speculative generation when the “tree confidence” drops below the threshold.
- Merge-based token tree construction
- 基础思想也很好理解:对于一个确定的输入,每个 SSM 都会并行地输出一个自己预测的序列,然后通过共享前缀的方式组合为 token tree.
- 这里有个问题:如何让不同的 SSM 趋向于生成不同的序列
- SpecInfer 采用了“迭代式专家微调”,利用一个"提示-序列"数据集,对第一个 SSM 进行微调,然后让每个后续的SSM都专注于学习前一个 SSM 预测失败的案例,来达到一种互补的效果。
Tree-based Parallel Decoding
Tree attention 部分感觉没什么新意,作者自己硬造了一个概念。
创新点集中在 token tree verification 这部分
把问题抽象出来,就是 kv cache 的存储和注意力计算都是线性的,但输入的 token tree 是树形的,怎么进行统一
- 对于 kv cache
- 让所有路径共用一个缓存,并通过 DFS 的顺序来更新这个缓存
- 对于 attention 计算
- tree 的性质:从根节点出发,有且仅有一条路径可以到达选择的节点,所以 SpecInfer 采用了 Topology-aware Causal Mask,屏蔽掉不应产生影响的Token之间的注意力计算(例如,兄弟节点之间),使得在一次大的矩阵运算中,就能并行且正确地完成树中所有节点的验证。
Multi-Step Speculative Sampling
对于需要随机性的“随机解码”(Stochastic Decoding),为了保证生成结果的概率分布与原始LLM完全一致,SpecInfer提出了一种名为MSS的验证算法。
工作原理: 与传统方法先从LLM采样一个Token再看是否在推测中不同,MSS会先检查推测的Token,并以一个精确计算的概率来决定是否接受它。 如果拒绝,它会根据剩余的概率分布重新调整,再对下一个推测的Token进行判断,直到接受一个Token或所有推测都被拒绝为止。
优势: 该算法被证明在数学上与原始LLM的输出分布完全等价 ,同时相比于其他采样验证方法,它拥有更高的推测接受率,能验证更多的Token。