Flow-Drafter-4B
A speculative-decoding draft model for Qwen/Qwen3.5-4B, from the Chained-Flow project.
Instead of drafting tokens one-at-a-time (a chain), it predicts all K future hidden states in a
single flow pass, expands them into a draft tree, and the base model verifies the whole tree in
one forward pass — accepting the longest path it agrees with (lossless).
Architecture — joint-VAE tree flow drafter
- A transformer-hidden VAE compresses the 2560-d backbone hidden into a 640-d latent; the
flow runs in that small latent (cheap, base-decoupled).
- A flow-matching drafter predicts K=8 future latents in 2 Euler steps, decoded back to hiddens.
- Token feedback at zero sequential cost: a low-rank Markov head (
bias = W2(W1[prev])) plus
path-conditioned residuals correct each branch by its chosen token.
- The VAE is unfrozen and trained jointly with the flow for acceptance (not reconstruction), with a
small reconstruction anchor — the recipe that made the latent flow work at scale.
Acceptance across model sizes (RedHatAI/speculator_benchmarks + held-out gsm8k)
Live tree-accept = real greedy tokens accepted per verify pass, measured against the backbone. All numbers are lossless (speculative output is identical to the base model).
| domain | Flow-Drafter-4B | Flow-Drafter-9B | Flow-Drafter-27B |
|---|
| gsm8k held-out (in-distribution) | 5.25 | 5.70 | 5.37 |
| math_reasoning (structured) | 5.44 | 5.25 | 5.73 |
| HumanEval (code) | 4.02 | 4.27 | 3.95 |
| writing (prose) | 3.28 | 3.55 | 3.48 |
| qa (short free-form) | 2.38 | 2.69 | 2.79 |
| summarization (prose) | 2.05 | 2.30 | 2.30 |
Accept tracks output predictability, not strict in/out-of-distribution: structured math/code stay high; free-form prose/QA/summarization dip. The vLLM decode speedup (compiled-forward + full-cudagraph verify + cudagraph draft) grows with base-model size because the latent-flow draft is base-decoupled — 9B measured 1.27× mean (1.77× math, 1.57× code).
Files
model.safetensors — the drafter (flow experts, Markov head, path head, and the jointly-trained VAE weights)chained_flow_tree_config.json — architecture + loss configvae/ — the base VAE checkpoint (needed to construct the VAE at load time)
Loading
Load with the chained-flow project's load_tree_module (set vae_dir to the bundled vae/ folder).
This drafter is meant to be driven as a tree proposer + lossless tree-verify over the Qwen/Qwen3.5-4B backbone.
Trained on a diverse mix (gsm8k, nemotron-math, nemotron-stem, alpaca-code, dolly-chat).
