Flow-Drafter-4B

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).

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.