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 config
• vae/ — 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).

Companion models: Flow-Drafter-9B, Flow-Drafter-27B.