<aside>

免责声明:本文依旧是人类手打,除画图外无 AI 参与

</aside>

0006-dflash2-banner.svg

<aside>

自回归的限制和甜点都是一体两面的:在给定的算力条件下,如何尽可能将并行计算发挥到极致,将串行计算压到最少,这是投机解码领域的一贯思路

</aside>

1. DFlash 的缺陷和优点

我们这里假设读者已经掌握了 DFlash 的初版思路:其草稿层权重并非自回归地,而是一个扩散模型,直接一步到位地生成出所有的草稿 Token。但是这样有个问题,没有体现出自回归模型的「依赖性」,也就是哪怕是草稿产生的Token,其内部的概率分布也不应该是一步到位的。这样的错配就是 DFlash 一代算法的致命缺陷: 接受率较低,生成的很多步 Token 最后都是浪费的。(当然,哪怕是浪费了许多,也是有用的,毕竟快啊)

从一代反思可得,最大缺陷还是局限在没有一个「串行」的过程,也就是强制将概率分布「矫正」回条件概率。于是后续的 Dspark 算法出现,将其进一步改进,接受率进一步升高。但是其马尔卡夫头的那部分开销依然是不小的。

2. 第一步:只要筛的足够多,接受率总会上去的

Metric 0 1 2 3 4 5 6 Acceptance length
Recall@1 85.4% 80.3% 79.4% 78.3% 77.5% 75.9% 72.9% 4.27
Recall@16 99.5% 97.3% 94.8% 92.6% 90.8% 89.4% 87.8% 6.79

这里是摘取自其博客中的一个对比表格。可以看到,他们团队观察到一个具体现象:诶,我每步只取 TOP1 的 Token 来做预测,接受率低。但是我要是每步取 TOP 16 个候选 Token,这接受率不是一下就上去了?

这也是个大实话,毕竟之前 EAGLE 算法也是整过 TOPK 这个参数的,那话又说回来,谁不知道力大砖飞呐,你一个位置上的预选 Token 就给我塞 16 个,你这步数一多,自回归特性一迭代,岂不是延迟起飞?我要你这预测头干嘛?

这里就是典型的想多了,或者说想少了。对 DFlash 来讲,它每一步保留多少个 Token 重要吗?只要我送进主模型做验证的时候,还是保持每个位置上只保留一个 Token 不就好了?

image.png

比如就像上图,只要我候选集里很大概率是有正确的Token 的,我只需要每一步提前替主模型算一下,最大概率选哪个不就好了?就像上图里,这个「Parallel Path Selector」组件,只要我计算过程足够轻巧,带来的 overhead 较小,哪怕力大砖飞的扩大候选集,还可以胜任这个工作。

那么 DFlash V2 是咋做的呢?这里又是借鉴了一下「Attention」的思路。我知道你很急,但是你别急,你看我说的对不对啊:我们已知,要对每一步的一堆候选词选出最合适的那个,且每一步都要合乎周礼,啊呸,不是,合乎逻辑,即第 N 步要和第 N-1 步保持逻辑关系,这不就是 Attention 的思路么?那么问题转化为,在计算量上考虑,哪种 Attention 开销最小?当然是我们的 MLA 大人,先降维,在一个较小的维度空间进行计算,岂不美哉?于是这里 DFlash V2 团队就巧妙借鉴了一下思路。按实际生产来设置参数,推理 7 步,候选池 16 个为例:

步骤 输入 输出 运算
取候选 logits (7,V) candidate_ids (7,16) + candidate_logits (7,16) 每个位置做 topk(16)
hidden States 降维 hidden states (7,2560) context features (7,256) H(h)=Linear(2560→256)
构造 predecessor ids anchor token (1) + candidate_ids (7,16) predecessor_ids (7,16) pos1 用 16 个 一样的anchor 重复;pos2~7 用前一个位置的 16 个候选
构造 successor ids candidate_ids (7,16) successor_ids (7,16) 当前位置自己的 Top-16 就是 successor
查 predecessor codebook predecessor_ids (7,16) + predecessor_codebook (V,256) A (7,16,256) 按 token id 去词表对应行找到对应的256 维的输出
查 successor codebook successor_ids (7,16) + successor_codebook (V,256) B (7,16,256) 按 token id 去词表对应行找到对应的256 维的输出
和 Hidden States 融合 A (7,16,256) + H(h) (7,256) (7,16,256) A * H(h)[:,None,:]
pair 打分 上一步 (7,16,256) + B (7,16,256) (7,16,16) 每个 predecessor × 每个 successor 做 256 维 dot product
加回原始 LM 分数避免前面的 pair 误差过大 pair scores (7,16,16) + candidate logits (7,16) (7,16,16) S(a,b)=pair(a,b)+logit(b)
选路径 7 个 score matrix 7 个 draft token 从 anchor 开始,逐位置 lookup / argmax

可以看到 DFlash2 Selector本质上是一个极轻量的、上下文感知的自注意力 QKᵀ 打分过程:它把 predecessor candidate 当作 Q ,把 successor candidate 当作 K,在 256 维低秩空间中一次性计算所有 16×16 路径打分;只是它不做 AV 信息聚合,而是直接把 score 用来选 token 路径。

3. 第二步:只要草稿层足够深,接受率总会上去的

众所周知,草稿层为啥预测不对 Token 呢?还不是拨给草稿层计算的算力和参数太少了,假如我草稿层直接来一手 10 多层的叠加计算,我这产生的预测 Token 正确率肯定很高!嗯,说的很好,DFlash 的团队也观察到了这一「显然的」趋势。

image.png

先不管那根红线,看看其他几个颜色的先,可以发现,层数越多,后续步数的接受率越高。相当于他们团队给我们验证了一下这个思路是确实可行的。那么现在的问题又来了:我们能不能「搏一搏单车变摩托」,拿 5 层去打 15 层呢?