跳到主要内容
返回时间线
bioRxiv来源发表:

EvSpark:面向混合DNA基础模型的无损推测解码

核心概要

EvSpark 为 Evo2 这类卷积-循环-注意力混合 DNA 基础模型设计了一套推测解码系统,通过一次块前向并行验证草稿并用“状态切片”选择保留的中间状态来同时恢复 FIR、IIR 与 KV 三类推理状态而无需重放,在 Evo2 7B 的 48 条提示基准上(三个训练种子)对 43 条真实序列提示取得 2.96× 加速、含五条合成对照时为 3.27×,在 262k 上下文仍保持 1.84×–2.43×,并在 20B/40B 目标上分别取得 2.18×–2.46×(真实序列)与 2.51×–2.78×(全量套件)。

Source-provided article image: EvSpark: Lossless Speculative Decoding for Hybrid DNA Foundation Models
Figure 1 ·

Figure 1: EvSpark architecture and speculative decoding. (a) Frozen Evo2 supplies intermediate features and offline teacher supervision; the 7B configuration extracts layer 27 (HCL). (b) A projected window of up to 64 target states forms the K/V prefix of each of two causal drafter layers (d = 1024). Anchor and mask embeddings produce all draft logits in one trunk pass; a previous-token Markov bias then supports serial sampling. Input and output embeddings are tied and frozen; projections and the two-layer trunk are trained. Dashed paths denote offline supervision using teacher- forced draft probabilities before the sampling transform T; serial sampling is shown for decoding. The confidence head does not control decoding. (c) The same Evo2 verifies the old anchor and six drafts in one block. Three drafts are accepted; a sampled correction becomes the next unconsumed anchor. (d) State returns to block index 3: shorten the readable KV prefix, select the last two outer-FIR inputs u2, u3, and retain IIR state s3. The feature buffer commits the same consumed prefix for the next round; no target computation is replayed.

bioRxiv · 第 4 页

深度剖析

提出针对混合目标的无重放验证:一次初始状态块前向加上统一的状态切片协议,可一致地恢复 FIR、IIR 与 KV 三类状态。 此前推测解码的状态恢复主要针对 Transformer 的 KV 缓存截断,或针对状态空间/混合模型的单独状态恢复;本文把 StripedHyena2 的有限卷积窗口、模态 IIR 状态与注意力缓存纳入同一份切片契约。 论文给出 Lemma 2(块前向契约)与 Lemma 3(状态恢复)及 Theorem 1(序列级精确性)的证明,并报告身份草稿对照在去掉重放后从 0.54× 提升到 1.05–1.16× 原生吞吐;切片本身约 0.62 ms,替代约 22 ms 的第二次目标前向。

设计了一个低延迟的并行、特征条件化蒸馏草稿器,用一次主干前向产出整块草稿。 该草稿器沿用 DSpark 式架构并适配到 DNA:以锚点与 W=64 的目标隐状态窗口为条件,两层宽度 d=1024 的因果注意力主干,配合全矩阵马尔可夫偏置与置信度头;论文在混合目标上评估了特征注入位置与训练预算。 草稿成本在测试长度上为 2.5–4.3 ms,速度从 1.77× 升到 3.07×;对比中 First-8 在 γ=12 时降到 0.85×、First-16 全程低于原生速度,独立 1B 自回归草稿器虽在 γ=12 达到 ¯τ=8.33(EvSpark 为 4.37),但单块草稿耗时 395 ms、仅 0.43×。

在规模与运行条件上验证了该实现与训练配方可迁移,并刻画了其运行边界。 把同一宽度-1024 草稿器配方迁移到 Evo2 20B 与 40B(分别取第 20/24 层与第 45/50 层隐状态),并给出长上下文、持续生成、批处理与多候选的测量。 20B/40B 在真实序列上 2.18×–2.46×、全量套件 2.51×–2.78×;草稿延迟在各目标上稳定在 3.6–4.0 ms,占原生单 token 步长的 9.3%(40B)与 18.4%(20B);262k 上下文下 E. coli 与 B. subtilis 为 1.84–2.43×;多候选 C=1/2/4 时 ¯τ 从 4.37 升到 5.40,但全量套件加速从 3.07× 降到 2.64×。

在调控 DNA 设计的完整工作流案例中,把解码加速转化为端到端设计时间缩短。 此前工作多报告单序列解码加速;本文在固定的预测器引导搜索流程下测量包含预填充、候选生成、缓存处理、打分与最终检查的完整设计时间。 在相同设备上分别校准出原生 batch 8 与 EvSpark batch 4,平均设计时间从 13.63 分钟降到 8.73 分钟,配对加速中位数 1.57×(范围 1.42×–1.70×);两个后端各四个输出均满足预设的预测资格标准(搜索集成与最终检查器 AUROC ≥0.90),平均检查器 AUROC 分别为 0.975 与 0.977。

启示与展望

该结果面向单序列解码延迟场景,token 生成基准不含预填充;调控设计案例则在固定搜索策略下额外覆盖预填充、批量独立候选与打分。加速在 7B 到 40B 目标、51k 至 262k 上下文以及 32k 持续生成上均有测量,但 262k 测试需要 48 GB 的 4090 变体,更大目标在单台 H20 上以 SDPA 评估。对希望降低长序列 DNA 生成或设计工作流延迟的读者,这套方法提供了可复用的状态切片契约与草稿器配方;对任意并发请求混合下的服务性能,论文明确留待后续评估。

精确性保证建立在精确算术之上,bf16 实现存在近并列的贪心情形与可测量的采样偏差,在重复序列中最明显,全模型 fp32 对照与序列级数值差异界仍待补充。预设的序列统计等价判据未被满足,事后分析给出的区间包含零但未建立等价。48 条提示覆盖区域有限、存在位点重叠与生物多样性代表性不均,解码轨迹变化可使套件均值移动最多 0.44×。20B/10M 单元未收到计划中的 ncRNA 窗口,40B/80M 的 hg38 被重复使用 1.50×,实际训练混合与计划存在差异。调控设计评估仅覆盖两种模式与两个设计种子,其质量与多样性描述的是观测输出,预测资格为计算评估,调控功能尚未经实验验证。

来源