Qwen3.8-Flash-Next DSpark 训练实验
以 Qwen3.8-Flash-Next 为冻结的 Target,训练三层 DSpark draft。各结构变体复用同一份离线数据与 hidden cache,比较其对收敛、候选选择和接受长度的影响。与 Qwen3-4B 消融 相比,本组使用不同的 Target、draft 深度和 aux 层,实验数值不作跨模型直接比较。
实验设置
各组实验采用下表所列基础训练设置,结构变体、附加目标与参数更新范围在各节分别说明。
| 项目 | 设置 |
|---|---|
| Target | Qwen3.8-Flash-Next;Target backbone、token embedding 与 LM head 冻结 |
| 训练数据 | PerfectBlend,经 Target 重新生成回复,由 DeepSpec 缓存特征 |
| 离线 cache | Aux hidden states 与 target_last_hidden_states;后者经冻结的 LM head 得到 teacher 分布 |
| 验证集 | GSM8K 256 条问题对应的重新生成回复及 hidden cache |
| 最大序列长度 | 4096 tokens,包含 prompt 与 response |
| Draft 架构 | 3 层;隐藏维度 2560;FFN 中间维度 9728;GQA:24 个 Q heads、2 个 KV heads,head dim 256 |
| Aux 层 | [45, 46, 47],按配置中的零基层号记录 |
| Block size | 7;Baseline 输入为 [anchor, MASK × 6] |
| Anchor 采样 | 每个样本最多 512 个 |
| 监督 | Response-only |
| Markov head | Vanilla,rank 256 |
| Batch size | 4 GPUs × micro-batch 4 × 梯度累积 32;global batch 512 |
| 学习率 | 峰值 \(3\times10^{-4}\);warmup 比例 0.04 |
| 训练计划 | 完整 scheduler 为 26,040 optimizer steps;约 2,604 步 / epoch |
| Baseline loss | \(0.1L_{\mathrm{CE}}+0.9L_{\mathrm{L1}}+L_{\mathrm{conf}}\);位置权重 \(\exp(-i/4)\),\(i=0,\ldots,6\) |
| 精度与随机种子 | BF16;seed 42 |
| Checkpoint | 本组对照实验每 651 步保存一次 |
学习率调度按完整的 26,040 步定义。一个 epoch 的阶段性实验在第 2,604 步停止,仍沿用完整调度的前段。对于从已有 checkpoint 初始化的实验,分别报告预训练步数与新增优化步数。
除冻结基线实验外,aux projector、draft backbone、Markov、confidence head 与各变体的新增参数联合优化。Target 相关参数和离线缓存保持固定。Engram 与 Draft Memory 沿用 baseline 损失,Prefix reranker 的附加排序目标和冻结策略在对应小节定义。
评估口径
离线指标
在同一参考前缀下,记 teacher 分布为 \(p_i\),draft 分布为 \(q_i\)。分布重叠为
每个有效 block 的概率接受长度定义为
不足 7 个有效位置时,无效位置的 \(a_i\) 置零,\(\tau_{\mathrm{prob}}\) 对有效 block 聚合。各位置的 accept_rate@i 按有效 token 数聚合,不使用损失的位置衰减权重。这些指标在缓存的参考前缀上以 teacher forcing 计算,区别于多轮推理中测得的 mean acceptance length(MAL)。
对 reranker,train/tau_probabilistic 和 eval/tau_probabilistic 使用重排后的分布:在 top-16 候选上对修正分数做 softmax,候选之外的概率置零。Teacher 仍使用原始全词表概率,不在候选集合内重新归一化。重排前的 DSpark 分布则在全词表上归一化。
| 指标 | 含义 |
|---|---|
rerank_tau_gain |
同一 checkpoint、同一批 anchor 上,重排后与重排前的 \(\tau_{\mathrm{prob}}\) 之差 |
rerank_target_agreement_gain |
重排前后选中 teacher argmax 的比例之差;0.02 表示增加 2 个百分点 |
reranker/candidate_recall |
Teacher argmax 落在 Markov top-16 内的比例 |
rerank_teacher_forced_prefix_gain |
与参考 token 序列连续匹配的前缀长度变化,仍在 teacher-forced 输入下计算 |
rerank_tau_gain 衡量同一模型内部的重排增量,包含残差修正、top-16 截断与重新归一化的共同影响。完整方法的收益则以相同训练预算下的独立 baseline 为参照,二者分别报告。
推理指标
推理指标包括真实 Target verification 下的 MAL、draft–verify cycle latency 和 tokens/s。对照条件固定为相同的测试集、采样方式、并发数、TP、最大长度和 proposal 数量。
当前 Engram 的实际前缀查表与服务链路尚未完成集成及在线评估,查表与 Target 计算重叠的调度也未验证。Draft Memory 尚未接入 vLLM 多轮推理。Prefix reranker 的采样接口仅支持 greedy,其随机采样收益不能由离线概率重叠推断。
实验分组
| 实验 | 初始化与更新范围 | 比较的问题 |
|---|---|---|
| Baseline | 从头训练原 DSpark | 原结构随训练预算增加的收敛水平 |
| Engram-Markov | 从头联合训练 draft 与 Engram 门控投影 | Markov head 的 n-gram 条件对接受长度的影响 |
| Engram-MASK:门控残差 | 从头联合训练 draft、输入投影与逐槽标量门 | 保留 MASK embedding 的门控注入效果 |
| Engram-MASK:直接替换 | 从头联合训练 draft 与输入投影 | 内容相关输入替代固定 MASK 的效果 |
| Draft Memory | 从头训练 draft 与 memory adapter,仅第二阶段回传 | 上一轮 draft 状态对新 anchor 下续写的作用 |
| Prefix reranker:联合训练 | 从头联合训练 draft 与 reranker | 前缀条件化重排相对独立 baseline 的收益 |
| Prefix reranker:冻结基线 | 加载 baseline step 2604,仅训练新加的 reranker | 固定 proposal 能力后,重排模块能提供多少增量 |
Baseline 架构
Baseline 将块内生成分解为并行特征计算与顺序 token 选择。三层 draft backbone 在已知前缀和 anchor 条件下,一次生成各预测位置的隐藏表示,LM head 将其映射为初始 logits。Vanilla Markov head 根据前驱 token 添加低秩词表偏置,confidence head 独立估计对应位置的分布重叠率。
图左以 D 为 anchor,展示 E–I 的生成,虚线表示相邻位置的选择依赖。右图展开预测 E 时的 Markov 修正与 confidence 分支。本文架构图按需省略输入槽,实际 block size 均为 7,包含一个 anchor 槽和六个后续输入槽,具体表示由各变体定义。
Target context 与并行预测
设 block size 为 \(B\),token embedding 为 \(E(\cdot)\)。对于前缀 [... A B C] 与新 anchor D,draft 输入为
MASK 槽共享输入 embedding,并使用不同的位置编号。输出采用 next-token 对齐:anchor 槽对应 \(h_1\),用于预测 E,后续槽依次预测 F–K。
对 anchor 之前的每个位置 \(j\),拼接 Target 第 45、46、47 层的隐藏状态,经投影与归一化得到 context 特征:
\(W_{\mathrm{aux}}\) 是可训练的无偏置投影。图中的 Target context 截止 C,anchor D 仅通过自身 embedding 输入,不提供其对应位置及未来位置的 Target hidden。
每层的 query 来自当前 draft 隐藏状态,K/V 由 context 特征与当前 block 的隐藏状态共同构成,在同一次 attention 中读取。三层共享 \(c_j\),分别使用自身的 K/V 投影,通过 attention、残差连接和 MLP 更新 draft 状态。最终经 RMSNorm 得到 \(h_1,\ldots,h_B\),再由冻结的 LM head 并行生成全词表初始 logits \(\ell_1^0,\ldots,\ell_B^0\)。
Block 内采用双向 attention,输入仅包含 anchor 和 MASK。训练可将多个 anchor block 合并到一次前向计算中,各 block 的可见范围限定为自身输入槽及对应 anchor 之前的 Target context。
Markov 顺序修正
令 \(t_0=D\),\(t_i\) 为第 \(i\) 个生成位置的 token。Markov head 将前驱 token 映射为低秩表示 \(m_i\),再投影为词表偏置 \(b_i\):
\(W_1\) 与 \(W_2\) 均参与训练。偏置 \(b_i\) 仅显式依赖前驱 token,不读取 \(h_i\),上下文信息由 \(\ell_i^0\) 提供。图右的第一步对应 \(m_1=W_1[D]\) 与 \(b_1=W_2m_1\)。Greedy 选择为 \(\hat t_i=\arg\max_v\ell_i(v)\),图示结果为 E。该修正作用于全词表,不限制候选集合。
训练采用 teacher forcing,所有位置的 Markov 修正可并行计算。推理使用实际选出的前驱 token,依次完成查表、logits 修正与 token 选择。图左的虚线仅连接相邻位置的选择过程,backbone 隐藏状态在整轮内保持不变。
Confidence head 与训练目标
Confidence 分支读取 draft 隐藏状态 \(h_i\) 与词表投影之前的 Markov 表示 \(m_i\),通过线性层与 sigmoid 得到预测值:
输入维度为 \(2560+256=2816\)。监督目标为同一参考前缀下 teacher 分布 \(p_i\) 与修正后 draft 分布 \(q_i\) 的重叠率 \(a_i=\sum_v\min(p_i(v),q_i(v))\)。该连续标签停止梯度,以二元交叉熵优化 confidence head。图右的 \(\hat a_1\) 是独立辅助输出,不参与 E 的选择,也不改变 CE、L1 的位置权重。
基础训练目标为 \(0.1L_{\mathrm{CE}}+0.9L_{\mathrm{L1}}+L_{\mathrm{conf}}\)。CE 监督记录的 next token,L1 对齐 teacher 与 draft 的全词表分布,confidence loss 拟合分布重叠率。三项损失均按有效 response 位置及位置衰减权重聚合。
离线 aux cache 提供模型输入,Target 最后一层的隐藏状态经冻结 LM head 生成 teacher 分布,仅用于监督与评估。
Engram 增强 Markov head
Engram-Markov 在前驱 token 的低秩表示中引入 n-gram 特征,并由当前 draft 隐藏状态控制注入强度。Backbone 与 MASK 输入保持不变。
图右展开以 D 为 anchor 预测 E 的第一步。Engram lookup 读取截至 D 的前缀 [... A B C D],返回特征 \(e(D)\),与 \(W_1[D]\) 融合后修正 logits。
门控与特征注入
从 1 开始编号生成位置,令 \(t_0\) 为 anchor,\(h_i\in\mathbb R^{2560}\) 和 \(\ell_i^0\) 分别为第 \(i\) 个预测位置的 draft hidden 与初始 logits。记 \(e_{i-1}\in\mathbb R^{2560}\) 为截至前驱 token \(t_{i-1}\) 的前缀查表得到的 raw Engram 特征,即 n-gram lookup 拼接后、尚未经过 K/V 投影的表示。
Draft hidden 提供 query,Engram 特征提供 key 和 value。记 \(N\) 为无可训练 affine 参数的 RMSNorm,Markov rank 为 \(r=256\),则
\(g_i\) 是每个位置一个标量,用来缩放该位置的 Engram value。这是单个 query–key 对的门控,不包含跨位置的 softmax attention。增强后的 Markov 表示和最终 logits 为
其中 \(W_1[t_{i-1}]\) 保留原有 unigram 路径,\(W_2\) 将融合后的 256 维表示映射为全词表偏置。Confidence head 同样读取增强后的 \(m_i\)。
三个无 bias 投影 \(W_Q,W_K,W_V\) 均为 \(2560\rightarrow256\),新增 \(3\times2560\times256=1{,}966{,}080\) 个 draft 参数。\(W_V\) 零初始化,使初始 Engram 残差为零。在 backbone 和原 Markov 参数相同的条件下,初始输出与 vanilla Markov 一致。
训练与推理
离线训练读取与参考序列对齐的 Engram 特征。预测 \(t_i\) 时使用前驱位置 \(t_{i-1}\) 的缓存,各位置的门控与 Markov 修正可并行计算。
推理时,lookup 随实际生成前缀逐 token 更新:读取 D 位置的特征预测 E,再基于包含 E 的前缀查表以预测 F。该路径不能沿用分叉后的参考序列缓存。上述参数量仅计入 draft 投影,Engram 表、特征传输和逐 token 查表构成额外存储与时延开销。
Engram 门控增强 MASK 输入
门控 Engram-MASK 在原 MASK embedding 上加入前缀特征残差,以各槽的可训练标量控制注入强度。门值零初始化,后续仍使用 vanilla Markov head。
图中 anchor 为 D,Engram 特征来自截至 C 的前缀 [... A B C]。特征投影后由各 MASK 槽共享,anchor 保留原 token embedding。
门控残差
令 \(a\) 为 anchor 在原序列中的位置,\(e_{a-1}\in\mathbb R^{2560}\) 为 pre-anchor 的 raw Engram 特征。\(P\in\mathbb R^{2560\times2560}\) 是无 bias 的可训练投影,RMSNorm 位于投影之后,且不含可训练 affine 参数。Block 输入为
\(\theta_j\) 是第 \(j\) 个相对 MASK 槽的可训练标量,跨样本和 block 共享。门值不随输入内容变化,特征依赖由 \(z_a\) 提供。\(\tanh(\theta_j)\in(-1,1)\) 允许正向或反向缩放残差,区别于 Engram-Markov 的内容相关门控。
所有 \(\theta_j\) 均零初始化,因此初始 \(u_{a,j}=E(\mathrm{MASK})\)。在原有 draft 权重相同的条件下,初始前向与 baseline 一致。投影 \(P\) 随机初始化,但门值为零时,它从任务损失获得的梯度也为零。门值离开零后,投影才开始接收该路径的任务梯度。新增参数为 \(2560^2+6=6{,}553{,}606\) 个,包括一个共享投影和六个标量门。
训练与推理
训练对每个 anchor 读取 \(a-1\) 位置的 Engram 缓存,缺少有效前驱或 block 无效时置零。六个槽共享投影结果,保留各自的位置编号、RoPE 与块内双向 attention。投影和标量门由基础任务损失学习,不增加独立门控目标。
推理每轮只读取一次 pre-anchor 特征,经投影与门控后输入 draft,整块生成期间保持不变。该特征在新 anchor 产生前即可确定,使 host lookup 具备与 Target 计算重叠的条件。
Engram 替换 MASK 输入
直接替换版本移除 MASK embedding 与标量门,以 pre-anchor Engram 的投影作为各槽输入,其余结构保持 baseline 配置。
图中以 D 为 anchor,读取截至 C 的特征 \(e(C)\)。D 通过原 token embedding 路径输入,后续槽由同一 Engram 投影填充。
输入替换
令 \(a\) 为 anchor 在原序列中的位置,\(e_{a-1}\in\mathbb R^{2560}\) 为该位置之前的 raw Engram 特征。一个无 bias 的线性投影 \(P\in\mathbb R^{2560\times2560}\) 将其映射到 draft 输入空间,随后进行无可训练 affine 参数的 RMSNorm:
六个槽共享 \(z_a\),位置编号为 \(a+1,\ldots,a+6\),沿用 RoPE 和块内双向 attention。
新增参数为投影矩阵 \(P\),共 \(2560^2=6{,}553{,}600\) 个。该投影随机初始化并参与训练,因此初始模型与 baseline 不等价。
训练与推理
特征选取与门控版一致:训练读取 \(a-1\) 位置的缓存,缺少有效前驱时置零。新增投影由基础任务损失直接优化。
推理每轮查表并投影一次,将 \(z_a\) 广播至所有替换槽,不随块内新选 token 更新。相较门控版,该路径省去逐槽缩放与残差相加,仍保留查表、特征传输、投影和归一化开销。
Draft Memory
Draft Memory 复用上一轮 draft 对后续位置的隐藏表示,为新 anchor 下的并行预测提供跨轮条件。Memory 来源于 draft,而非 Target 在被拒绝候选路径上的状态。训练采用共享参数的两阶段前向过程,复用已有 Target hidden cache,无需额外 Target rollout 或 Engram 特征。
两阶段训练
第一阶段在无 memory 条件下执行不保留梯度的 draft 前向计算,取得最终 RMSNorm 后的隐藏表示。LM head 与 Markov head 在参考前驱 token 条件下计算 argmax,并确定其与记录序列的首次不匹配位置。在此位置之前,参考前驱与 greedy draft 已选 token 一致,因此该过程与自回归 greedy 生成得到相同的首次参考错误。后续预测不参与 anchor 转移。
图中旧 anchor 为 C,Target context 截止 B,参考续写为 D、E、F、G、H。若前两个预测匹配,第三个预测为 X,则以参考 token F 为新 anchor,将 context 扩展至 E。对应已匹配位置及首次不匹配位置的 \(h_1,h_2,h_3\) 被移除,仅保留 \(h_4,h_5\),供第二阶段预测 G、H、I、J、K。
该转移依据记录序列模拟拒绝,而非实时 Target verification。Memory 在 Markov 选择之前由旧 anchor 与 MASK 并行计算,不包含沿错误 token X 重新计算的状态,也不附加候选 token embedding 或拒绝标记。
第二阶段以新 anchor、对应的正确 Target context 和筛选后的 memory 为输入,采用 teacher forcing 计算基础任务损失。梯度仅通过该阶段回传,联合更新共享 draft 与 memory adapter。
Memory 筛选与读取
令旧 anchor 的绝对位置为 \(a\),block size 为 \(B\),第 \(i\) 个输出 hidden 为 \(h_i^{\mathrm{old}}\),其中 \(i=1,\ldots,B\)。该 hidden 的旧 query 位置为 \(a+i-1\),预测的 token 位置为 \(a+i\)。若连续匹配了 \(m<B\) 个 token,则
新 anchor 为 \(a'\)。后缀索引集合 \(\mathcal I\) 中的隐藏表示经共享的无偏置线性投影和带可训练尺度参数的 RMSNorm 转换为 memory:
\(P\in\mathbb R^{2560\times2560}\) 随机初始化,\(\gamma\) 为归一化尺度参数。Adapter 共增加 \(2560^2+2560=6{,}556{,}160\) 个参数,不引入门控。停止梯度作用于第一阶段的隐藏表示,\(P\) 与 \(\gamma\) 由第二阶段的损失优化。
Memory 的 RoPE 位置采用预测位置 \(a+i\),即旧 query 位置编号加一。因此,首条保留状态对齐新 anchor 后的第一个预测位置 \(a'+1\)。图中 \(h_4,h_5\) 分别对齐 G、H。
Adapter 只计算一次,得到的同一组 \(M\) 供三层 draft 使用。每一层的 query 仍来自当前 draft states,K/V 则由正确 context、memory 和当前 draft states 拼接后,通过该层原有的投影得到。省略归一化与 RoPE 后,可写为
其中 \(H_{\mathrm{ctx}}\) 为投影后的正确前缀特征,\(X_\ell\) 为当前层的 draft states。Memory 作为静态 K/V 来源,与 context 和当前 block 共享同一次 attention 归一化,不增加独立 cross-attention。可见性约束限定为本 anchor 的有效 memory、当前 draft block 与 anchor 之前的正确 context,排除其他 block 及未来 Target states。
若整块全部匹配,则转移至记录序列的 bonus 位置,memory 为空。缺少后续监督或跨越回答边界的转移不参与训练。本组实验保留全部有效后缀,不使用 memory dropout。
训练与推理
相对 baseline,训练增加一次无梯度 draft 前向计算、首次错误定位所需的输出头计算,以及 memory 投影与扩展 attention 的开销。虽然不增加 Target 前向计算,相同步数仍不对应相同计算预算,训练效率需结合实际耗时评估。
在线推理需按请求保存上一轮 draft 的最终隐藏表示。Target verification 后,根据接受长度移除已接受状态及首次拒绝状态,将剩余后缀按预测位置对齐至下一轮。各层通过自身投影生成 memory K/V。
训练与推理仍存在两项分布差异:拒绝位置由参考序列而非实时 Target 验证确定,训练 memory 仅来自无 memory 的第一阶段,而连续推理中的源隐藏状态可能已包含更早一轮 memory 的影响。
Prefix reranker
设计动机
Vanilla Markov 的偏置只显式依赖单个前驱 token,缺少对当前 block 内更长已选前缀的编码。
Prefix reranker 通过块内因果 Transformer 编码 anchor 与已选 token,再结合当前 draft 隐藏状态,对 Markov top-16 候选进行残差打分。
图右展开第三个预测位置:前缀 [D, E, F] 编码为 \(r_3\),与 \(h_3\) 和候选特征共同生成修正分数。图示候选 G 在残差修正后超过 H。
候选与前缀表示
设 \(t_0\) 为 anchor,\(t_i\) 为第 \(i\) 个生成位置的 token。训练时使用参考 token,推理时使用当前轮已选 token。记 backbone 输出为 \(h_i\in\mathbb{R}^{2560}\),冻结 LM head 生成的初始 logits 为 \(\ell_i^0\)。
复用 Markov head 的两张秩为 256 的 token 表 \(E_{\mathrm{in}}\) 与 \(E_{\mathrm{out}}\),得到原 DSpark 分数
候选集合由 Markov 修正后的分数确定。因此,teacher argmax 的候选覆盖率构成重排选择准确率的上界。
Prefix Transformer 的输入为 \([t_0,\ldots,t_{i-1}]\)。各 token 经共享的 \(E_{\mathrm{in}}\) 查表、线性投影并加上可学习的块内位置 embedding。一个 pre-norm 因果 Transformer 层生成前缀表示 \(r_i\)。其宽度为 256,包含 4 个 attention heads,FFN 为 \(256\rightarrow1024\rightarrow256\),使用 SiLU,输出经过 LayerNorm。
前缀编码范围限于当前 block,历史 context 由 \(h_i\) 提供。
残差打分
将 \(h_i\) 投影到 256 维,与 \(r_i\) 拼接,再经过 MLP 和输出投影形成查询向量 \(u_i\)。候选 token 则从共享的 \(E_{\mathrm{out}}\) 查表并投影为 \(k_v\):
最终选择为 \(\hat t_i=\arg\max_{v\in C_i}s_i(v)\)。Prefix Transformer 编码前缀一次,候选通过查询向量与投影特征的内积并行评分。
\(W_o\) 零初始化,并在并列分数下优先保留原 DSpark 的选择,从而保证初始 greedy 选择一致。候选内归一化与全词表归一化仍有差异,因此这一初始化不保证概率重叠指标相同。
训练目标与更新范围
训练沿用现有 hidden cache。每个 block 的前驱输入右移一位,以 anchor 开头,其后为参考序列中的 token。Prefix Transformer 使用 causal mask,在一次 forward 中并行编码所有位置。
记 \(p_i\) 为 cached reference prefix 上的 teacher 分布,重排监督标签为
该标签可能不同于参考序列中实际记录的 token。记 \(\mathcal H\) 为有效且 \(v_i^\star\in C_i\) 的位置集合,只在这些位置计算候选内的排序 CE:
其中 \(w_i=\exp(-(i-1)/4)\),求和覆盖 batch 中各 block 的有效命中位置。没有命中位置时,该项为零。联合训练保留重排前的原始 DSpark 目标:
原目标仍监督全部有效位置,包括 teacher argmax 未进入 top-16 的位置。联合训练中,排序损失也会更新 backbone 和共享 Markov 参数,因此内部的“重排前模型”不等同于独立训练的 baseline。
冻结基线实验以训练 2,604 步的 baseline 为起点,固定 backbone、Markov 和 confidence head,仅优化 reranker 的新增参数,并重新初始化优化器与学习率调度。曲线横轴采用累计训练步数:预训练步数加上 reranker 的新增优化步数。横轴偏移不影响后者的学习率调度。
推理开销
Backbone 每轮只运行一次。第 \(i\) 个位置利用上一个已选 token 更新 Prefix Transformer 的本地 K/V,得到 \(r_i\),并计算 Markov 修正后的 top-16 候选及其残差分数。选出的 token 进入下一位置,整块候选最后交给 Target 验证。
Prefix K/V 在当前 block 内增长,每轮重新清空。新增计算位于串行选择路径,包括前缀编码、候选筛选与残差打分,不增加 Target rollout。
结果总览
实验结果尚待按统一评估口径汇总。