让快速推测生成更准，通过修正分布偏差提升效率
TreeFlash: Parallel AR-Approximation for Faster Speculative Decoding
- 用MLP结合隐藏状态和前一个词，逼近自回归分布
- 在多种模型上实现12%块效率提升、9%加速比
- 适合追求高速生成的推理系统开发者
一次生成块的推测解码方法通过单次前向传播生成完整草案，极大提升了吞吐量。然而，其每个草案词仅基于前缀上下文预测，不依赖已生成词，导致草案分布随生成深度增长逐渐偏离验证器的真实自回归分布。这一问题在树状生成中尤为严重，不同分支被迫共享后续词的相同边缘分布。我们提出TreeFlash，通过引入一个条件于草案隐藏状态和前一词的MLP层，近似自回归分布。TreeFlash采用两阶段近似机制，在保持一次生成器$/mathcal{O}(1)$解码复杂度的同时，解决了分布偏移问题。在多个任务与模型上，TreeFlash表现达到领先水平，相较边际树生成方法，块效率提升12%，速度提升9%。
原文摘要 · Abstract (English)
One-shot block drafters for speculative decoding generate the full draft in a single forward pass, achieving strong throughput by eliminating sequential token generation. However, they predict each draft token conditioned only on the prefix context, with no dependence on previously drafted tokens. This non-autoregressive conditioning causes the drafter's distribution to diverge from the verifier's true autoregressive distribution as draft depth grows. This problem becomes more severe in tree-based drafting, where distinct branches are forced to share the same marginal distribution for subsequent tokens. We propose TreeFlash, which addresses this by incorporating an MLP layer conditioned on the drafter's hidden state and the previous token to approximate an autoregressive distribution. TreeFlash retains the $\mathcal{O}(1)$ decoding time complexity of one-shot drafters by employing a two-stage approximation mechanism. TreeFlash achieves state-of-the-art performance across a variety of tasks and models, improving over marginal tree drafting by $12\%$ higher block efficiency and $9\%$ higher speedup.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。