论文
RheoSampling:解决随机动态树推测解码中的 One-Hot 困境
RheoSampling: Resolving the One-Hot Dilemma in Stochastic Dynamic-Tree Speculative Decoding
摘要
推测性解码通过并行起草多个标记来加速 LLM 推理,基于树的方法通过层次结构进一步提高效率。诸如 EAGLE-3 之类的动态树方法通过确定性的 top-K 扩展和全局剪枝在贪婪解码下表现良好。然而,在随机解码(T>0)中,这种机制将草稿分布折叠成单热概率,导致接受率严重下降。这就造成了一个困境:动态树方法牺牲随机采样来保留上下文感知拓扑,而静态树方法则保留具有上下文无关结构的随机采样。出现这个问题是因为相同的概率分布用于两个相互冲突的任务:构建树和验证词元。由于产生的随机过程,这种耦合使得随机性的直接注入具有挑战性。我们通过解耦这些角色来解决这个问题:RheoSampling 为从草稿分布中采样的词元分配一个用于树扩展和修剪的代理概率以及用于验证的真实采样概率。具体来说,我们在确定性的 top-K 个槽中注入一个采样词元,并在构建和验证过程中以不同的概率对其进行处理,使 RheoSampling 成为第一个具有上下文感知的 top-K 构建和随机采样的动态树方法,同时保持无损。我们通过等价类分析建立无损保证,将随机树空间压缩为易于处理的类。基于OT的验证策略和稀疏草案机制确保理论成果转化为实践效率。大语言模型和基准测试的实验表明,与最先进的动态树方法相比,接受率和加速速度都有所提高。该框架可以提供用于分析随机树结构的模板。