Rethinking the Roles of Diffusion and Decoding: Predictive-State Flows for Parallel Language Generation
Organizations: Georgia Institute of Technology
Abstract
In continuous diffusion language models (DLMs), latents generated by the diffusion prior are mapped almost deterministically to tokens, either by a decoder or through a token embedding matrix. Unlike in image generation, where visual details can vary continuously, the diffusion prior must capture fine-grained discrete token choices in continuous space, which can make the joint latent distribution difficult to model. We argue that generative roles should instead be divided: the diffusion prior should generate high-level predictive information, while the decoder should resolve the remaining token-level choices. Based on this view, we introduce Predictive-State Flows (PSF), which partitions a sequence into blocks and associates each block with a predictive state computed solely from the preceding context. A rectified-flow prior generates the sequence of predictive states, after which a block-causal decoder generates all blocks in parallel, autoregressively only within each block. This design reduces diffusion from L token-level states to K << L predictive states and reduces token-level sequential depth to L/K. On unconditional OpenWebText generation, a 197M-parameter model with K=8 achieves a generative perplexity of 22.8 at entropy 5.16 using 0.31 TFLOPs per 1024-token document, versus 24.0 and 7.2 TFLOPs for ELF-B with 32 sampling steps. The corresponding measured operating points yield MAUVE@512 of 0.93 and 0.73, respectively. Controlled comparisons further show that content states, which additionally encode their own blocks, yield 8-16% higher generative perplexity than predictive states, and that replacing sampled predictive-state trajectories with a shared mean lowers MAUVE@512 by 0.11. These results suggest that latent design for continuous DLMs should consider not only how compactly text is represented, but also how generative roles are divided between the prior and the decoder.
Figures & tables
| Model | Size | TFLOPs | GPU time (ms/doc) | Gen. PPL | Entropy | MAUVE |
|---|---|---|---|---|---|---|
| Held-out OpenWebText | — | — | — | 16.9 | 5.41 | — |
| GPT-2 small (AR) | 124M | 0.27 | 128.6 | 12.09 | 5.122 | 0.75 |
| AR, same data | 197M | 0.43 | 219.2 | 15.37 | 5.151 | 0.95 |
| MDLM | 170M | 347 | 16 386.7 † | 35.37 | 4.897 | 0.44 |
| BD3-LM, block 4 | 170M | 17.9 | 1 416.5 † | 23.75 | 5.039 | 0.61 |
| ELF-B, 32 steps | 105M | 7.2 | 20.9 | 23.97 | 5.158 | 0.73 |
Appendix figures & tables21 assets
Supplementary material from the paper’s appendix.
Appendix
| Network | Layers | Width | Heads | Params | Attention and positions | Inputs and outputs |
|---|---|---|---|---|---|---|
| Encoder | 6 | 512 / 1024 | 8 | 35M / 109M | block-causal self-attention; rotary embeddings with a hybrid within-block / across-block scheme | per-block mean pooling; linear head to the code posterior , |
| Block decoder | 18 | 512 / 1024 | 8 | 109M / 369M | causal self-attention within a block; cross-attention to with a causal mask over blocks; rotary embeddings within blocks | added at every position: -conditioned memory, and (zero-initialized linear layers), the source label (zero-initialized 2-dimensional embedding), and from position 2 the landing embedding of |
| Prior | 12 / 24 | 768 / 1024 | 12 / 16 | 88M / 308M | full self-attention over tokens | tokens of dimension through a 256-dimensional input bottleneck; 4 guidance-scale tokens and 4 time tokens |
| Stage | Steps (k) | Batch | KL budget | Tokens (B) | |
| PSF, Stages I–II | |||||
| base (Stage I) a | 30 | 512 | 64 | — | 14.6 |
| landing target b | 10 | 512 | 64 | — | 4.9 |
| full recipe c | 30 | 512 | 64 | — | 14.6 |
| self-realized prefixes d | 10 | 512 | 64 | — | 4.9 |
| realization codes e | 10 | 512 | 64 | 16 | 4.9 |
| Published | Re-evaluated | ||||
|---|---|---|---|---|---|
| Model | Sampling | Gen. PPL | Entropy | Gen. PPL | Entropy |
| ELF-B | SDE, 32 steps, guidance 3, | 24.1 | 5.15 | 23.97 | 5.158 |
| ELF-B | SDE, 64 steps, guidance 3, | — | — | 19.42 | 5.096 |
| ELF-M | SDE, 64 steps, guidance 3, | 21.7 | 5.18 | 21.65 | 5.184 |
| ELF-L | SDE, 64 steps, guidance 3, | 23.3 | 5.28 | 23.65 | 5.281 |
| Cosmos, | Euler SDE, 200 steps | — | — | 90.65 | 4.760 |
| Steps | Gen. PPL | Entropy | Loop | |
|---|---|---|---|---|
| 32 | 1.5 | 23.97 | 5.158 | 4.0% |
| 40 | 1.5 | 21.80 | 5.138 | 6.6% |
| 48 | 1.5 | 19.22 | 5.097 | 14.6% |
| 64 | 1.5 | 17.57 | 5.067 | 13.1% |
| 64 | 1.0 | 19.42 | 5.096 | 14.6% |
| 32 | 0.5 | 39.40 | 5.296 | 1.6% |
| Setting | Gen. PPL | Entropy | Loop |
|---|---|---|---|
| 9.80 | 4.855 | 10.2% | |
| 14.32 | 5.051 | 5.9% | |
| 17.72 | 5.137 | 5.4% | |
| 21.81 | 5.233 | 4.7% | |
| , nucleus 0.85 | 12.90 | 5.082 | 4.1% |
| , nucleus 0.86 | 13.26 | 5.057 | 5.7% |
| Model | Setting | Seeds | Mean sd | Reference |
| PSF | plain, at entropy 5.16 | 22.5 / 22.5 / 23.2 | 22.8 | |
| PSF-L | plain, at entropy 5.16 | 18.4 / 18.6 / 18.6 | 18.4 | |
| PSF-L | plain, at entropy 5.16 | three seeds | 20.3 | |
| PSF | , guidance 4 (measured) | four seeds | (entropy ) | 24.80 |
| PSF | nucleus, zero codes at inference, at 5.16 | 19.1 / 19.4 / 19.5 | 19.3 (Figure 10 ) | |
| PSF | states-only, nucleus 0.6, at 5.16 | 19.4 / 19.4 / 19.3 | 19.4 |
| Comparison | Gen. PPL at matched entropy | Diagnostic |
|---|---|---|
| Landing target added | seam perplexity 9241 958 | |
| No landing target (final model) | to | codes kept |
| Realization code added | to (45.4 40.6 at 5.16) | greedy repeats 55% 14% |
| Code lookahead added | to | block-final tokens to |
| Self-realized prefixes added | (69.8 65.3 at 5.34) | true plan without label |
| Source label added | (65.3 63.9 at 5.34) | restores true-plan decode |
| Variant | Sampler | Gen. PPL | vs. control | Loop |
|---|---|---|---|---|
| PSF | ||||
| shuffled codes | plain | 23.4 at 5.174 | (23.7) | 9.1% (7%) |
| codes | plain | 19.95 at 5.164 | (23.0) | 19.3% |
| true plan, true zero codes | 19.9 at 5.168 | (24.5) | 18.1% (14%) | |
| retrained w/o code channel | nuc., ban | 19.0 | 0% | |
| states-only prior | nuc. 0.55, ban | 19.0 | 0% | |
| Teacher-forced NLL | Gen. PPL @5.16 | True-plan decode @5.16 | Loop | ||||||
|---|---|---|---|---|---|---|---|---|---|
| PSF | PSF-L | PSF | PSF-L | PSF | PSF-L | PSF | PSF-L | ||
| 64 | 16 | 2.63 | 2.45 | 39.2 | 42.2 | 40.3 | 35.1 | 1.2% | 7% |
| 32 | 32 | 2.78 | 2.57 | 32.2 | 27.1 | 31.0 | 25.6 | 2.8% | 9% |
| 16 | 64 | 2.89 | 2.69 | 24.9 | 20.3 | 26.5 | 21.2 | 7.6% | 10% |
| 8 | 128 | 2.97 | 2.77 | 22.8 | 18.4 | 24.0 | 19.1 | 6.9% | 11.2% |
| Model (sampler) | Train FLOPs | FLOPs/doc | Serial | Latency | Throughput (doc/s) | |
| ( ) | (TFLOPs) | depth | (batch 1, s) | batch 32 | batch 256 | |
| PSF ( ) | 4.2 | 0.80 | 66 | 0.14 | 62.6 | 71.4 |
| PSF ( ) | 4.6 | 0.52 | 82 | 0.18 | 70.2 | 92.3 |
| PSF ( ) | 4.9 | 0.38 | 114 | 0.26 | 60.4 | 93.7 |
| PSF , plain ( ) | 5.2 | 0.31 | 178 | 0.42 | 36.7 | 62.9 |
| PSF , states-only ( , ban) | 5.2 | 0.31 | 178 | 0.45 | 30.8 | 48.9 |
| GPT-2 Large, at entropy | Qwen2.5, at 5.16 | Window | |||||
| System | 5.00 | 5.10 | 5.16 | 5.25 | 1.5B | 7B | at 4.845 |
| PSF | 28.5 | 34.5 | 39.2 | 48.4 | 35.0 | 32.3 | 38.1 |
| PSF | 23.4 | 28.3 | 32.2 | 40.6 | 28.0 | 25.7 | 26.7 (extr.) |
| PSF | 17.7 | 21.8 | 24.9 | 31.4 | 22.0 | 20.2 | 22.2 |
| PSF | 15.9 | 19.7 | 22.8 | 27.64 (5.237) | 20.1 | 18.4 | 20.4 |
| states-only, | — | — | 19.4 | — | 16.7 | 15.1 | 17.1 |
| System | Distinct-4 | Self-BLEU-4 | 13-gram overlap |
|---|---|---|---|
| Held-out OpenWebText | 0.905 | 0.148 | 2.67% |
| PSF , | 0.803 | 0.301 | 0.54% |
| PSF-L , | 0.781 | 0.296 | 0.83% |
| PSF , states-only, nucleus 0.55, ban | 0.750 | 0.404 | 0.49% |
| PSF-L , states-only, nucleus 0.6 | 0.741 | 0.363 | 0.96% |
| ELF-B, 32 steps | 0.801 | 0.292 | 0.04% |
| Sampler | Entropy | PPL with ban | PPL without ban | PPL (%) | Loop, no ban | |
| PSF , prior’s codes | ||||||
| 4 | 5.111 | 20.25 | 20.2 | 9% | ||
| 4 | 5.162 | 22.76 | 22.9 | 6.9% | ||
| 3 | 5.235 | 26.25 | 27.5 | — | ||
| 6 | 5.120 | 20.50 | 20.6 | — | ||
| 8 | 5.109 | 19.82 | 20.1 | — | ||
| Qwen2.5-7B, words 1–300 | Qwen2.5-32B coherence, words | ||||
| System 1 vs. system 2 | Coherent | Human-written | 1–300 | 300–600 | all |
| real vs. real (calibration) | 65:42 | 41:41 | — | — | — |
| PSF plain (prior’s codes) vs. real | 6:143 | 12:131 | 1:198 | 0:195 | 1:196 |
| PSF states-only + ban vs. real | 1:166 | 2:147 | 0:194 | 0:197 | 0:196 |
| PSF-L states-only + ban vs. real | 4:144 | 10:147 | — | — | — |
| ELF-B, 32 steps vs. real | 10:127 | 7:121 | 0:193 | 0:194 | 1:196 |
Text Generation
Text generation studies methods for producing coherent, useful language from prompts, data, and dialogue context.
See all 422 papers