Lafa: 8B flow LM on <30B tokens, 2.54× AR speed via self-speculation

Large Flow Language Models

Arnav Chavan, Chanhyuk Lee, Udbhav Bamba, Jaehoon Yoo, Seunghoon Hong, Aditi Raghunathan, Jinwoo Kim

cs.CL, cs.AI

2026-10-21

Lafa adapts pretrained AR models into blockwise asynchronous flow LMs, up to 8B on <30B tokens. At 7B, self-speculative decoding reaches 2.54× AR throughput at similar quality.

What problem this solves

An autoregressive model emits one token per forward pass. Teacher forcing makes training parallel. Generation stays serial, and that loop is where latency sits.

Three lines of work try to commit more tokens per pass. Multi-token prediction (MTP) guesses extra tokens from the same prefix. Discrete diffusion masks tokens and fills them back iteratively. A continuous flow language model embeds each token and integrates a path from Gaussian noise to the clean embedding. Because the state is continuous, few-step jumps and guidance transfer from image and video models. Discrete diffusion has no equivalent interface.

Earlier flow language models stayed small and were shown on simple tasks. At a billion parameters, training efficiency failed first. Lafa starts from an existing autoregressive checkpoint, uses under 30B tokens, reaches 7B and 8B, and still caches the prefix from block to block.

Method

Text is generated in blocks, left to right, with block length B set to 4, 8, 16, or 32. The next block sees only tokens already committed. Prefix keys and values ignore noise inside the block, so every denoising step and every later block reuse one cache.

Inside the block the model does flow matching on d-dimensional word embeddings. Noise is Gaussian. The endpoint is the clean embedding. A standard flow shares a single time t across the block. Lafa assigns each position its own time τ, and positions on the left denoise first. τ is not hand-designed. It is the normalized inverse CDF of that position's denoising loss, estimated online. High-loss noise levels move slowly, so uniform draws of t put more samples and more integration steps there. The learned schedule runs faster on the left.

The network predicts a distribution over the clean token. Velocity comes from an identity: the direction from the noisy embedding toward the predicted embedding, multiplied by that position's time derivative. Cross-entropy supervises the class. Squared error supervises the direction.

Weights are initialized from a pretrained autoregressive model: Qwen2.5 Instruct at 0.5B and 7B, Qwen3 8B Base at 8B. Time is injected with adaLN, initialized as the identity. Prefix positions stay at time 1. Their attention is causal and cannot see the noisy block. Attention inside the block is bidirectional, so a partly cleaned token on the right can revise the left. Prefix outputs still predict the next token, and that head later verifies drafts.

Embeddings change direction only. Their length is locked. A growing norm would dilute a fixed noise scale.

Training adds three terms: prefix cross-entropy, denoising cross-entropy, and embedding squared error. The clean sequence and the noisy blocks share one forward. Self-conditioning reuses the previous prediction, so inference does not add a forward. Prefix guidance, in the same family as classifier-free guidance, evaluates the block once more without the prefix.

A flow map distills the integration into a jump from time s to t. On the diagonal the student matches the teacher with KL. Off the diagonal it matches two shorter jumps. The teacher's prefix predictions are distilled as well. One-step generation is the jump from 0 to 1.

Self-speculation denoises once from pure noise and drafts a whole block. The autoregressive head keeps the longest prefix that agrees with greedy decoding, left to right. If the first token disagrees, it commits the autoregressive token. The string equals that model's greedy decode. Verifying one block and drafting the next happen in the same forward.

Results

Training runs on Dolci-Instruct-SFT for 5 to 15 epochs, about 9B to 27B tokens. That is at least 18× fewer tokens than the published budgets of Dream, LLaDA, and Efficient-DLM at similar scale, which go up to about 500B. Baselines marked † were retrained from the same initialization, data, and budget. The headline score is a macro average of IFEval, BBH, MMLU-Pro, GPQA, GSM8K, and HumanEval. Generation is capped at 1024 tokens, 512 on GSM8K and 640 on HumanEval.

Single-token accuracy:

Method0.5B7B8B
AR SFT†28.758.456.8
MTP†25.953.9n/a
DeepSeek-MTP†30.260.0n/a
MARS†30.459.7n/a
Dreamn/a55.6n/a
LLaDAn/an/a43.5
Efficient-DLMn/an/a63.5
Lafa†31.660.963.1

Lafa beats matched AR SFT at every scale, by 2.9, 2.5, and 6.3 points. At 8B, 63.1 sits next to Efficient-DLM at 63.5, with an order of magnitude less data.

Flow decoding improves as Euler steps increase. With B=4 and enough steps, 0.5B approaches AR SFT and 7B passes it. Larger blocks trade accuracy for parallelism. The flow map, distilled at B=4, beats the teacher at one and two steps. At 7B the two-step student approaches the four-step teacher, leads at four steps, and ties around eight. With still more steps the student trails the undistilled flow.

Self-speculative text equals greedy decoding of that checkpoint. The accuracy gain is the autoregressive head improving under joint training. TPF is how many tokens one forward commits.

SetupAvg accuracyAvg TPF
AR SFT 7B58.41.0
Lafa 7B, B=460.92.6
Lafa 7B, B=859.34.3
Lafa 7B, B=1658.04.9
AR SFT 8B56.81.0
Lafa 8B, B=463.12.9

Wall-clock throughput is measured on one H100 at batch size 1, averaged over IFEval, GSM8K, and HumanEval. Lafa 7B at B=16 reaches 2.54× the throughput of AR SFT, which the paper calls comparable quality. At B=4 it is about twice as fast and more accurate on that average. The plotted range is 2.02× to 2.54×. DeepSeek-MTP and MARS match or beat AR SFT quality and reach at most 1.4× its throughput. Discrete diffusion sits below AR SFT, and faster settings lose more.

At 0.5B, B=8 yields about twice the TPF of DeepSeek-MTP, MARS, and PTP, with a higher average. At 7B, B=8 matches MARS at B=4 on accuracy and reaches about 1.5× its TPF.

Ablations use 0.5B, B=4, and 32 Euler steps, on the same three-task average. Bidirectional attention inside the block scores 38.0, fully causal attention 32.7. Asynchronous time scores 38.0, synchronous time 36.2. Norm-locked embeddings score 38.0, unconstrained embeddings 36.9, frozen embeddings 34.4. Denoising cross-entropy alone is 33.6, embedding squared error alone 29.7, both together 35.8, and both plus the prefix autoregressive loss 38.0. Self-conditioning moves the score from 32.3 to 38.1. Prefix guidance moves it from 34.3 to 38.1.

Why it matters

The fast path you can deploy is self-speculation. The few-step flow map has not reached the quality of a long integration. One set of weights drafts and verifies. There is no second drafter and no extra MTP head. The prefix cache stays. The data bill is an instruction-tuning run, under 30B tokens, against published discrete-diffusion adaptations that use up to about 500B. Joint training raises average autoregressive accuracy.

Multi-step flow still needs enough Euler steps before it matches AR SFT. The 2.54× result is the greedy autoregressive string, with more than one token accepted per forward.

Limitations

The paper says the flow map still trails flows that take more steps, and that self-conditioning and prefix guidance are not distilled cleanly. One-step drafts disagree across positions. Larger blocks cost accuracy. Accepting drafts by a confidence threshold costs more accuracy, which is why the default remains self-speculation.

The 2.54× number is one H100 at batch size 1. Speculative decoding usually pays less once a large batch fills the GPU, and that regime is not measured. The output matches greedy decoding token for token.

Averages hide the task split. At 8B, against the authors' own Base SFT, IFEval moves from 38.1 to 62.3, GSM8K from 62.9 to 77.8, and GPQA from 29.5 to 31.7. MMLU-Pro stays at 56.2. BBH moves from 71.6 to 70.7, and HumanEval from 82.3 to 79.9. Most of the 6.3-point gain is IFEval and GSM8K. The 2.54× setting is B=16. On the six-task average that checkpoint scores 58.0, below AR SFT at 58.4, and HumanEval falls from 78.7 to 73.8. B=4 raises the six-task average to 60.9, while HumanEval is still 76.8.

The 8B AR SFT is not the public Qwen3 Instruct checkpoint. Further training of Instruct got worse, so both sides start from Base. Against Efficient-DLM the scores are 63.1 and 63.5, on different data. The paper does not test editing or infilling, so the control argument for continuous flow is unchecked. Each block length is a separately trained checkpoint.

Terms

Source

What people are saying

Related papers

All paper explainers