# 并行草稿模型的最佳因果修正方案

- 来源：Hao AI Lab (@haoailab)
- 发布时间：2026-07-31 04:23
- AIHOT 分数：51
- AIHOT 链接：https://aihot.virxact.com/items/cms7yv1lx02noro20krh27avf
- 原文链接：https://x.com/haoailab/status/2082925179529994586

## AI 摘要

Sky Computing Lab 通过训练所有骨干网络与因果修正设计的组合，系统研究了并行草稿模型的最佳因果修正方案。DSpark 的 Markov head（约 78M 参数）在 logits 形成后添加低秩偏置，而 JetSpec++ 的注意力层（约 75.5M 参数）在隐藏空间修正。两种方案均采用条件树构建，沿分支顺序运行修正后再对子节点评分。

## 正文

http://x.com/i/article/2082909485799198720

# What Makes a Good Causal Correction for Parallel Drafting

Authors： @aaronzhfeng · @Lanxiang_Hu · @YuYangQian_ai · @Jensen_Yuan · @haozhangml

In our last post we set out where DSpark and JetSpec agree and where they differ. Both lines arrived at the same conclusion independently： once drafting is cheap， what limits a parallel drafter is whether its proposals stay coherent as a sequence， and the fix is a causal correction， which conditions each drafted position on the tokens preceding it along its own branch rather than on its slot in the block.

What neither line has done is ask what the best causal correction looks like， which is the question this post takes up. In this blog， we trained every combination of the backbone and causal correction designs. We train all models including DSpark on that same data rather than directly using DeepSeek's checkpoint for a fair comparison.

## 1. Causal Correction Design Choices

A block-parallel drafter has two parts， and they raise different design questions.

The backbone drafts a whole block in one forward， with KV injection reading the target model's cached attention states， and both lines build on the same design. JetSpec's block mask is causal： the draft head applies a tree-causal mask under which each candidate node attends to the prefix and to its own ancestors， never to descendants or to unrelated siblings， so every branch carries the target's own left-to-right dependency while all nodes are still scored in a single pass.

The causal correction is the newer question， and the two lines diverge on where in the forward pass it lands：

- After the logits are formed. DSpark's Markov head adds a low-rank bias to the output distribution， conditioned on the immediately preceding token， at roughly 78M parameters.

- Before the logits are formed. JetSpec refines in hidden space， and this is where we introduce JetSpec++， which adds an attention layer to it： an attention block with the feed-forward half left out， revising the hidden state at each node from the tokens actually realized along that branch， at 75.5M parameters.

The two sit on either side of the vocabulary projection， lm_head， which turns a hidden state into a score for every token： one is the last operation before it， the other the first operation on the distribution， and that placement governs both what the correction is allowed to read and what a tree costs to build.

Conditional Tree Construction

Chain drafting compares the two as they already are. Tree drafting， the ground JetSpec was designed for， needed something built first.

The tree takes one distribution per draft position from a single backbone forward and expands the top-k of each， so every node at a given depth offers the same children whatever branch it sits on. That is how block-parallel drafters built trees before conditioning entered the picture， and it is where a tree of that kind leaks： tokens that are individually plausible compose into branches that do not hold together， and the verifier throws the whole branch away.

The construction algorithm is the same for all tree inference settings in this post： a best-first walk over accumulated log-probability， expanding the most promising node， scoring its children， repeating until the budget is spent.

The only change is： every expanded node is corrected under its own ancestry using the correction layer （with Markov head or attention layer） before its children are scored. Since a node's children cannot be scored until its parent is chosen， the correction runs sequentially along each branch. We call this conditional tree construction.

What differs is where the node scores come from. The same search runs on either distribution：

- A first-order transition. DSpark's Markov head shifts the backbone's logits at each node by a bias read from the parent token alone.

- Path-conditioned logits. JetSpec++'s attention layer revises the node's hidden state under its whole ancestry， and the revised state is projected into fresh logits.

> Note： We standardize our experiment under the setting： Qwen3-8B on a single B200， greedy， at most 1024 new tokens， n=100 per benchmark and AIME25's full 30. Trees run to a 255-node budget at maximum width 7 and chains at width 1， with accepted length τ counted per verification round， including the verifier's one bonus token.

## 2. DSpark vs. JetSpec++

Architectures and Training Configurations

The tables show： at depth 7 the two are nearly tied， separated by about a tenth of a token； at depth 15 JetSpec++ leads on every benchmark in both modes， by 0.61 to 0.92 accepted tokens in tree drafting and 0.24 to 0.52 in chain， widest on HumanEval in chain. The AIME25 margins sit in the noisiest column here， n=30 against everyone else's 100. The deeper you draft， the more the correction's design matters， which is precisely the regime a tree exists to exploit.

Performance at Matched Parameter Sizes

At essentially the same size， 75.5M against roughly 78M， the attention layer accepts 0.89 more tokens per round in tree drafting and 0.54 more in chain. Tripling it to 226.5M by restoring the MLP， the third row， adds 0.01. Parameter counts are plainly not what separates these designs. What a parameter is separates them.

- The Markov head's parameters are storage. Every correction is identified based on a lookup table shaped at training time， and the correction only touches the logits from the backbone.

- The attention layer's parameters are a machine. Every correction is computed on the fly using causal attention that provides conditioning on the whole path rather than the single position used in Markov head.

Attention Layer Input Designs

The layer is block-local and causal， and it is a residual adapter initialised to the identity， so it can only add what it learns. We trained three versions of it， differing only in what goes in.

The first row is the control： a full attention layer with the parent token withheld lands below every other row. Swap its input to the parent token alone and it gains 0.59 in tree drafting and 0.39 in chain； fusing the hidden state back in adds another 0.02 and 0.09. Nearly all of the correction's value at this depth is carried by the parent token， and the fusion's share grows with depth： at depth 15 the learned fusion is worth about 0.21 more in tree drafting than a plain sum of the two inputs.

## 3. JetSpec++ Extrapolates to Deeper Tree Depth

A draft head is trained at a fixed block size， but at inference time if a draft head can extrapolate to deeper depths， it will have a higher cap for speedup. The two attention patterns behave very differently when it is widened.

That is a reason to expect the two to come apart when asked for more depth， and it costs nothing to check： a head trained for 15 draft positions can be asked for 23 with no retraining and no new weights. We run the ablation and the results are the following：

At depth 23， the trained-depth margins of 0.71 in tree and 0.40 in chain become 1.67 and 1.24， a larger separation than either design produced at the depth it was actually trained for. JetSpec++ gains on all four benchmarks in both modes； DSpark， in tree drafting， gains on one and loses on three， and in chain it loses on all four.

To find the limit， we ran the same checkpoint at deeper positions.

The free depth is not unlimited， and it is nearly all spent by 23. Getting there is worth +1.66 accepted tokens in tree drafting and +0.88 in chain； the eight positions after it add 0.11 more in either mode. Whatever a served system does with this， the useful range is the first stretch.

## 4. Where the Capability Comes from？

To separate the two axes the campaign covers every combination： backbone （causal or bidirectional）， crossed with correction （an attention layer or a Markov head）. We run two additional ablation experiments and we summarize results in the table below：

From the table， we can draw the conclusion： causal backbones provide token acceptance at deeper tree depths than both bidirectional ones， and their results are on-par at trained depths. For causal correction， attention layer in general performs better than Markov head.