arXiv cs.LG精选82

长上下文起草不需要长上下文起草器

Long-Context Drafting Does Not Require Long-Context Drafters

· 1小时前

精选理由

证明无需训练超长上下文草稿模型,也能让推测解码在百万级上下文保持稳定加速。

AI 摘要

arXiv:2610.20865v1 公告类型:新 摘要:推测解码将廉价的起草器与针对目标模型的并行验证结合起来。在超长上下文中,这种方法不再划算,因为对完整前缀进行注意力计算的起草器会继承目标模型依赖上下文的成本,因此它原本旨在提供的加速会先逐渐减弱,随后反转。通常的应对方式是构建长上下文起草器:通过训练或微调一个起草器,或使用目标模型检查点中自带的多 token 预测头。我们认为这并无必要,因为现有起草器都以目标模型为条件,读取目标模型在整个前缀上计算出的隐藏状态;因此,远处的上下文无需经过起草器自身的注意力计算就已经传递给它们,而验证过程会纠正它们遗漏的内容。基于此,我们冻结已发布的短上下文起草器,并通过两种推理时机制让它们远超训练长度运行:一种限制其可见的注意力缓存,从而固定其成本;另一种将其旋转位置索引重新映射到训练范围内。整个过程无需改动任何权重,唯一需要调节的参数很早就达到饱和,因此只需一项 vLLM 侧配置即可覆盖从 4k 到两百万个 token 的长度,在短上下文下的加速损失最多为 0.16 倍。我们证明了在何种条件下,统一的索引偏移不会改变起草器的计算结果,并将接受率分解为两个可测量且具有严格界限的差距。无论冻结起草器在短上下文中带来多少加速,在三组目标模型—起草器组合和五个语料库上,这种加速都能保持到两百万个 token。该论断关注的是加速的平坦性,而非其倍数;移除任一机制都会使这种平坦性消失。该收益的前提是,每个生成 token 的推测成本始终低于匹配的仅使用目标模型的成本。

查看原文(arXiv cs.LG)