想象一下,你正试图写一个故事,主角是一位非常聪明但动作缓慢的图书管理员(即 AI 模型)。每当你向他询问故事的下一个词时,图书管理员都必须停下来,苦苦思索,查阅他那庞大的整座图书馆,然后向你低声耳语出下一个词。这就是当前 AI 的工作方式:一次一个词,一步一个词。 它很准确,但很慢。
这篇论文介绍了一种让这位图书管理员变得更快且不失准确性的新方法。他们称之为 “来自边缘分布的树”(或 DFlash-TfM)。以下是其工作原理的拆解,采用了简单的类比。
问题:“猜谜游戏”的限制
为了提高速度,研究人员发明了一个技巧,叫做 投机采样(Speculative Decoding)。
- 旧方法: 一个快速的初级助手(“草拟者”)猜测接下来的几个词。然后,缓慢的图书管理员(“验证者”)检查这些猜测是否正确。如果正确,图书管理员会一次性接受所有猜测。如果不正确,图书管理员会修正错误并重新开始。
- “因子化”草拟者的缺陷: 一些助手之所以超级快,是因为它们会一次性猜测接下来的所有单词,而忽略了这些词之间的相互联系。这就像一位厨师在还没品尝前一个食材的情况下,就直接猜出了汤里的下三个食材。
- 代价是: 随着猜测列表变长,厨师的猜测准确度会下降。第一个猜测可能是对的,但第三个通常是错的,因为它没有考虑到前两个词。这限制了每次可以被接受的单词数量。
解决方案:“编织者”助手
作者创造了一个结合了快速厨师的速度与谨慎编辑逻辑的新系统。他们将这个新编辑器称为 Weaver(编织者)。
- “Top-K”候选名单: 首先,快速助手(DFlash)进行快速、粗略的猜测,并为下一个位置提供一份包含 512 个最可能单词的候选名单。这就像一位厨师说:“我觉得下一个食材很可能就在这 512 种香料之中。”
- 编织者的任务: 编织者不再盲目猜测,而是观察这份候选名单。它扮演着一位聪明的编辑角色,会说:“好吧,如果第一个词是‘盐’,那么下一个词几乎肯定是‘胡椒’,而不是‘糖’。”
- 构建树结构: 编织者不仅仅是做出一条直线式的猜测。它构建了一棵 树。
- 想象一棵家谱树。根部是当前的句子。
- 编织者向外分支,创造出故事的不同可能路径(例如,“猫坐在垫子上” vs “猫坐在地板上”)。
- 因为编织者规模很小,且仅观察快速助手提供的候选名单,所以它构建这棵可能性之树的速度极快。
验证:检查这棵树
现在,缓慢的图书管理员需要检查这棵猜测树。
- 旧有的问题: 如果图书管理员使用标准的“循环”记忆系统(例如现代 AI 中的 Gated Delta Net 层),检查一棵树通常是一场噩梦。这就像试图通过逐一走遍树上的每一条分叉来查看哪条路径是真实的。这非常缓慢。
- 新的窍门: 作者发明了一种特殊的数学捷径(一种“无回溯”算法)。
- 他们没有逐一走遍每个分支,而是使用了一种 掩码三角求解(masked triangular solve)。你可以把它想象成一张神奇的地图,让你能同时观察整个树状结构,并瞬间知道哪条路径是正确的,而无需为每一个分支都重新计算记忆状态。
- 这就像拥有一个 GPS,它能瞬间在复杂的地图上标出正确的路线,而不需要你先去开遍每一条死胡同。
结果:速度与效率
通过结合这些想法,该系统实现了两大胜利:
- 接受更多的单词: 因为编织者修正了快速助手的逻辑错误,图书管理员接受的单词链变得更长(比之前的最佳方法多出高达 77%)。
- 巨大的加速: 整个过程极其高效,AI 生成文本的速度比标准慢速方法快了 4.37 倍。它还比之前的“最快”方法提升了约 25%。
总结类比
- 标准 AI: 一只正在写故事的蜗牛,一次写一个字母,并对照字典检查每一个字母。
- 旧的快速方法: 一个速读员试图猜出整个段落,但因为没注意到开头,经常把中间的部分猜错。
- 这种新方法(Weaver): 一个速读员快速挑选出 500 个可能合适的词,然后由一位微型且超级聪明的编辑(编织者)将这些词瞬间排列成逻辑严密的树状句子。随后,一个特殊的“神奇地图”(新的内核)会瞬间检查整棵树,以确定哪条路径才是真实的。
结果是,这种 AI 既拥有速读员般的写作速度,又具备细心编辑般的准确度,使交互过程感觉更加即时且响应迅速。
技术摘要:来自边际分布的树 (DFlash-TfM)
问题陈述
自回归大语言模型 (LLMs) 本质上是序列化的,每步仅生成一个 token。这种序列化特性造成了内存带宽受限的瓶颈,使得延迟主要由 DRAM 带宽而非计算量决定,从而导致能量效率低下且交互性受限(每序列生成的 token 速率)。虽然投机解码 (speculative decoding) 通过使用快速“草拟 (drafter)”模型来提出多个 token,供主“验证 (verifier)”模型并行验证,提供了一种解决方案,但现有方法面临着权衡:
- 分解式草拟器 (Factorized Drafters): 像 DFlash 这样的模型将未来的 token 预测为在并行状态下条件独立的边际分布 (marginals)。虽然高效,但其独立性假设会导致随着投机预算(草拟长度)增加,接受率急剧下降,因为它们无法捕捉所提议 token 之间的条件依赖关系。
- 自回归草拟器 (Autoregressive Drafters): 这些模型能正确捕捉依赖关系,但在每一步都会进行全词表投影,从而产生了高昂的开销,抵消了速度增益。
- 在循环模型上的验证 (Verification on Recurrent Models): 标准的树验证算法依赖于注意力掩码 (attention masks)。然而,使用门控 Delta Net (GDN) 层(例如 Qwen3.6)的现代模型利用了非交换状态转移,这类模型不支持简单的祖先限制掩码,这使得高效的单次遍历树验证成为了一个开放性问题。
方法论:DFlash-TfM
作者提出了 DFlash-TfM (Trees from Marginals),这是一种混合架构,结合了分解式草拟器的并行效率与自回归模型的条件准确性,并配备了一种新型系统级验证内核。
1. Weaver 适配器 (The Weaver Adapter)
核心创新在于 Weaver,这是一个轻量级的自回归适配器(56.7M 参数),运行在受限的词表之上。
- 输入: Weaver 接收来自预训练的分解式草拟器 (DFlash) 和验证器的隐藏状态。
- 机制: Weaver 不会对全词表进行投影,而是提取由 DFlash 预测的 Top-K 边际 token。随后,它在这一小组候选集(K=512)上自回归地预测残差分布。
- 收益: 这恢复了 token 之间的条件依赖关系(例如,token t+1 取决于已实现的 token t),同时避免了全词表投影带来的内存带宽成本。该模型构建了一个提案树,其分支是从这些经过条件化的残差中采样的。
2. 树构建 (Tree Construction)
作者修改了 DySpec 树构建算法以提高吞吐量。他们不再采用严格的顺序最佳优先扩展,而是从候选堆中提取前 w 个节点并同时进行 w 次扩展。这平衡了内存带宽与计算压力,使得构建完整的草拟树仅需 ⌈B/w⌉ 次顺序操作。
3. 面向 GDN 的无回滚树验证 (Rollback-Free Tree Verification for GDN)
为了支持具有 Gated Delta Net 层的模型,作者推导出了一种无回滚树验证算法:
- 挑战: GDN 层通过 St=αt(I−βtktkt⊤)St−1+βtktvt⊤ 更新状态。这种矩阵更新的非交换性质阻止了用于对角状态空间模型的标准累积乘积优化。
- 解决方案: 作者利用了线性递归的分块形式 (chunked form)。他们将验证建模为在树结构上的掩码三角求解 (masked triangular solve)。
- 树按拓扑顺序排序。
- 算法计算代表祖先关系的交互矩阵(X 和 Y)。
- 它通过求解线性方程组 (I+X)U=βV 来计算输出,而无需特意更新循环状态。
- 状态提交被延迟到下一次解码迭代,即在确定的接受分支已知之后。这消除了状态回滚的需求。
- 实现: 这在 SGLang 中实现为一个融合的 CUDA 内核,它计算 Gram 矩阵,通过分块前向替换求解系统,并在单次遍历中应用输出阶段。
核心贡献
- 混合架构: 引入了 DFlash-TfM,它使用轻量级自回归适配器 (Weaver) 来对分解式边际分布进行条件化,从而提升了纯分解式草拟器在长深度下的接受率上限。
- 理论洞察: 分析表明,分解式草拟器的接受率受限于边际分布与条件分布之间的全变分 (Total Variation, TV) 距离,并且通过 Weaver 对已实现 token 进行条件化,混合模型可以超越任何仅基于边际分布的草拟器的理论极限。
- 系统优化: 推导出了适用于 Gated Delta Net 层的新型无回滚树验证算法,使得对于之前难以优化的模型(如 Qwen3.6)进行高效的投机解码成为可能。
- 实现: 为 SGLang 开发了优化的 CUDA 内核,通过融合验证步骤,显著降低了相比于逐分支循环基准的延迟。
实验结果
在 NVIDIA B200 GPU 上针对 Qwen3.6-27B (bfloat16) 在对话、数学和代码基准测试上进行了评估:
- 加速比: DFlash-TfM 相比标准自回归解码实现了 4.37× 的加速。
- 基准对比: 它比高度优化的 DFlash 基准提升了 24.7%,比 DDTree 基准在不同数据集上提升了 10–30%。
- 接受长度: 与链式 DFlash 基准相比,该方法将平均接受长度 (MAL) 提升了 77%;在相同树规模下,比 DDTree 提升了 32%。
- 验证效率: 融合内核的扩展性明显优于逐分支循环验证。在树规模为 128 个 token 时,融合内核比循环基准快 7.1 倍。验证步骤仅占总解码时间的约 12%。
意义与主张
论文声称,分解式草拟器的主要限制是结构性的(独立性假设),而非缺乏模型容量。通过引入一个在受限词表上运行的轻量级自回归适配器,作者证明了在不产生全词表投影开销的情况下,恢复条件依赖关系并实现高接受率是可行的。
此外,这项工作解决了关键的系统差距:使具有非对角循环层 (GDN) 的模型能够进行高效的投机解码。提出的无回滚验证内核消除了状态回滚的需求,使得基于树的投机解码能够适用于更广泛的现代架构。作者总结道,这些结合了模型与系统的贡献,代表了在本地和服务器端 LLM 推理交互性方面的最先进水平。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。