Rethinking the Roles of Diffusion and Decoding: Predictive-State Flows for Parallel Language Generation

Jinjin He; Xinsong Feng; Lucheng Fu; Sinan Wang; Bo Zhu; Jiayun Wang

cs.CLOct 9, 2026

Rethinking the Roles of Diffusion and Decoding: Predictive-State Flows for Parallel Language Generation

Authors: Jinjin He, Xinsong Feng, Lucheng Fu, Sinan Wang, Bo Zhu, Jiayun Wang

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

Panel of Figure 1 — Figure 1: Content latents and predictive states. (a) A content latent encodes the current block, carrying predictive information and specific token choices. (b) In PSF, ${\bm{s}}_{k}$ encodes only the prefix before block $k$ ; the prior models predictive states jointly and the decoder predicts local realizations. Small cells schematically indicate predictive information (blue) and token choices (tan). Landing targets and optional realization codes are omitted; Appendix A.1 gives details.

Model	Size	TFLOPs	GPU time $\downarrow$ (ms/doc)	Gen. PPL $\downarrow$	Entropy	MAUVE $\uparrow$
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	$\approx$ 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

Table 1: Measured operating points (settings: Appendix A.7 ; matched entropy: Figure 2 ). H100 time: ms/document at batch 32 (

\dagger

: eager; Appendix C ).

Figure 2: Quality per FLOP and per second at entropy 5.16; lowest-perplexity setting per PSF size and

K

. (b) Lines connect batch 32 (small dots) and 256 (larger markers); ELF labels give documents/s.

\dagger

: own entropy. Sampling and interpolation: Appendix A.7 .

Figure 3: Perplexity–entropy comparisons. (a)

K{=}8

plain (solid) and states-only nucleus (hollow, dotted), without the 4-gram ban. ELF-B: 32/64 steps; BD3-LM: block 4. Hollow dashed: baseline guidance/temperature sweeps. (b) Landing-target ablations at

K\in\{8,64\}

; colors as in Table 3 .

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 7 — Figure 7: Parallel block decoding. Branch $k$ reads the cumulative memory $\hat{{\bm{s}}}_{\leq k}$ and generates its block autoregressively. All branches share parameters and advance synchronously for $\ell=L/K$ steps. The dashed path shows a landing target: after the first step, the sampled first token of the next block (“fox” here) conditions the preceding branch from step 2 onward.

Network	Layers	Width	Heads	Params	Attention and positions	Inputs and outputs
Encoder $E_{\psi}$	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 $({\bm{\mu}}_{k},\log\bm{\sigma}_{k})$ , $d_{r}=64$
Block decoder $p_{\theta}$	18	512 / 1024	8	109M / 369M	causal self-attention within a block; cross-attention to ${\bm{s}}_{\leq k}$ with a causal mask over blocks; rotary embeddings within blocks	added at every position: ${\bm{s}}_{k}$ -conditioned memory, ${\bm{r}}_{k}$ and ${\bm{r}}_{k+1}$ (zero-initialized linear layers), the source label $\lambda_{k}$ (zero-initialized 2-dimensional embedding), and from position 2 the landing embedding of $x^{(k+1)}_{1}$
Prior $p_{\phi}$	12 / 24	768 / 1024	12 / 16	88M / 308M	full self-attention over $K$ tokens	tokens of dimension $d+d_{r}$ through a 256-dimensional input bottleneck; 4 guidance-scale tokens and 4 time tokens

Appendix

Table 4: Architecture of the three networks. Widths and parameter counts are given as PSF / PSF-L (

d=512

1024

Stage	Steps (k)	Batch	$K$	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

Appendix

Table 5: Training schedule of the encoder and block decoder (Stages I–II of Section 2.4 ) and the prior (Stage III). Batch = documents per step (1024-token canvas); KL budget = code budget in nats per block (

\ell

nats, i.e. one nat per token; “—” = no code channel yet); tokens = steps

\times

batch

\times

953 non-padding tokens. Notes below the table say what each stage adds and when the encoder is frozen.

		Published		Re-evaluated
Model	Sampling	Gen. PPL	Entropy	Gen. PPL	Entropy
ELF-B	SDE, 32 steps, guidance 3, $\gamma{=}1.5$	24.1	5.15	23.97	5.158
ELF-B	SDE, 64 steps, guidance 3, $\gamma{=}1.0$	—	—	19.42	5.096
ELF-M	SDE, 64 steps, guidance 3, $\gamma{=}1.0$	21.7	5.18	21.65	5.184
ELF-L	SDE, 64 steps, guidance 3, $\gamma{=}1.0$	23.3	5.28	23.65	5.281
Cosmos, $L{=}512$	Euler SDE, 200 steps	—	—	90.65	4.760

Appendix

Table 7: Baselines: sampling settings, published generative perplexity and entropy, and the same quantities under our protocol (1000 documents, GPT-2 Large judge, first document of each canvas). The BD3-LM rows at nucleus 0.95 and 1.0 complete the sweep behind the log-linear interpolation at entropy 5.16 used in Section 3.1 and Figure 2 : 42.1 / 52.4 / 58.8 for blocks 4 / 8 / 16. Checkpoints: ELF-{B,M,L}-owt-torch , bd3lm-owt-block_size{4,8,16} , mdlm-owt , the ar.ckpt released with MDLM, and gpt2 . ELF and MDLM run in bfloat16; BD3-LM uses its released FP32 backbone; GPT-2 small and our models use TF32. a whole canvas; b 300 samples; c whole canvas under our judge: 38.2 / 5.28; d whole canvas under our judge: 12.8 / 5.28.

Steps	$\gamma$	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%

Appendix

Table 8: ELF-B (105M) with 32–64 SDE steps, 1000 documents, our judge; guidance 3 unless given in parentheses.

Setting	Gen. PPL	Entropy	Loop
$t{=}0.80$	9.80	4.855	10.2%
$t{=}0.85$	14.32	5.051	5.9%
$t{=}0.88$	17.72	5.137	5.4%
$t{=}0.90$	21.81	5.233	4.7%
$t{=}1$ , nucleus 0.85	12.90	5.082	4.1%
$t{=}1$ , nucleus 0.86	13.26	5.057	5.7%

Appendix

Table 9: The same-data AR baseline under the paper’s protocol (1000 documents, GPT-2 Large judge, per-document entropy, loop rate). Nucleus 0.86 and 0.88 were sampled in a later run than 0.85 and 0.90, hence the non-monotone entropies.

Figure 8: Perplexity–entropy sweeps by block count. Plain-sampler curves for PSF and PSF-L at each

K

: color denotes block count, solid lines with filled circles denote PSF, and dashed lines with hollow circles denote PSF-L. Gray diamonds mark ELF-B at 32 and 64 steps. The vertical reference is entropy 5.16.

Figure 9: Sampling choices and repetition at

K{=}8

. Left: PSF; right: PSF-L. Top: Gen. PPL; bottom: loop rate. Temperature sweeps fix nucleus 0.95; nucleus sweeps fix temperature 1. Both rows show the same 36 settings with loop rate

\leq 20\%

(1000 documents each); native

t{=}p{=}1

settings are omitted. The vertical line marks entropy 5.16.

Model	Setting	Seeds	Mean $\pm$ sd	Reference
PSF $K{=}8$	plain, at entropy 5.16	22.5 / 22.5 / 23.2	$22.7\pm 0.4$	22.8
PSF-L $K{=}8$	plain, at entropy 5.16	18.4 / 18.6 / 18.6	$18.5\pm 0.1$	18.4
PSF-L $K{=}16$	plain, at entropy 5.16	three seeds	$20.5\pm 0.2$	20.3
PSF $K{=}16$	$t{=}0.70$ , guidance 4 (measured)	four seeds	$24.81\pm 0.07$ (entropy $5.166\pm 0.003$ )	24.80
PSF $K{=}8$	nucleus, zero codes at inference, at 5.16	19.1 / 19.4 / 19.5	$19.3\pm 0.2$	19.3 (Figure 10 )
PSF $K{=}8$	states-only, nucleus 0.6, at 5.16	19.4 / 19.4 / 19.3	$19.4\pm 0.1$	19.4

Appendix

Table 10: Seed-to-seed variance of the interpolated or measured generative perplexity (1000 documents per seed, GPT-2 Large judge). “Reference” uses the original evaluation sweep at the entropy specified in each row.

Comparison	$\Delta$ Gen. PPL at matched entropy	Diagnostic
Landing target added	$-25\%$	seam perplexity 9241 $\to$ 958
No landing target (final model)	$+13$ to $+15\%$	codes kept
Realization code added	$-7$ to $-11\%$ (45.4 $\to$ 40.6 at 5.16)	greedy repeats 55% $\to$ 14%
Code lookahead ${\bm{r}}_{k+1}$ added	$-2$ to $-3\%$	block-final tokens $-6$ to $-7\%$
Self-realized prefixes added	$-6.4\%$ (69.8 $\to$ 65.3 at 5.34)	true plan $+30\%$ without label
Source label $\lambda_{k}$ added	$-2.1\%$ (65.3 $\to$ 63.9 at 5.34)	restores true-plan decode

Appendix

Table 11: Recipe ablations at

K{=}64

, compared at matched entropy. Added components are measured against the model without them when introduced; removal from the final model is marked separately (seed sd 0.04–0.4, Appendix A.8 ).

Figure 10: Realization-code ablations. (a) Change in generative perplexity after replacing the prior’s codes with zeros, with entropy matched within each labeled

K

row. (b)

K{=}8

, entropy 5.16; dashed guides mark the full-model references.

\pm

: standard deviation over three seeds. Full/zero

K{=}8

results also appear in Table 4 .

PSF $K{=}8$
Variant	Sampler	Gen. PPL	vs. control	Loop
shuffled codes	plain	23.4 at 5.174	$-1\%$ (23.7)	9.1% (7%)
codes $\sim\mathcal{N}(0,I)$	plain	19.95 at 5.164	$-13\%$ (23.0)	19.3%
true plan, true $\to$ zero codes	$t{=}0.65$	19.9 at 5.168	$-19\%$ (24.5)	18.1% (14%)
retrained w/o code channel	nuc., ban	19.0	$-10\%$	0%
states-only prior	nuc. 0.55, ban	19.0	$-10\%$	0%

Appendix

Table 12: Additional realization-code and retraining controls (

K{=}8

, entropy 5.16 unless stated). “vs.” uses the full model at matched sampler/entropy, with its PPL/loop in parentheses; content states use the corresponding predictive-state control. nuc. = nucleus; ban = 4-gram ban; ⋆ : Table 4 ; ‡ : extrapolated by 0.015 nats (zero codes) / 0.025 nats (true plan), Appendix B.2 .

		Teacher-forced NLL		Gen. PPL @5.16		True-plan decode @5.16		Loop
$K$	$\ell$	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%

Appendix

Table 13: Reported numbers behind Figure 4 : block length and scale with the recipe held fixed (plain sampler, the prior’s codes at every

K

). Teacher-forced NLL = held-out cross-entropy of the block decoder in nats per token; perplexities at entropy 5.16; true-plan decode as defined in Appendix A.3 ; loop rate at the setting closest to 5.16.

Model (sampler)	Train FLOPs	FLOPs/doc	Serial	Latency	Throughput (doc/s)
	( $\times 10^{19}$ )	(TFLOPs)	depth	(batch 1, s)	batch 32	batch 256
PSF $K{=}64$ ( $t{=}0.7$ )	4.2	0.80	66	0.14	62.6	71.4
PSF $K{=}32$ ( $t{=}0.7$ )	4.6	0.52	82	0.18	70.2	92.3
PSF $K{=}16$ ( $t{=}0.7$ )	4.9	0.38	114	0.26	60.4	93.7
PSF $K{=}8$ , plain ( $t{=}0.75$ )	5.2	0.31	178	0.42	36.7	62.9
PSF $K{=}8$ , states-only ( $p{=}0.55$ , ban)	5.2	0.31	178	0.45	30.8	48.9

Appendix

Table 14: Cost of one 1024-token document on one H100. FLOPs from PyTorch’s operator counter over the sampling code (batch 1; analytic for the same-data AR); train FLOPs

=6ND

(ours cumulative along the chain, rollouts and prior included; ELF-B 46B tokens, assumed for ELF-M/L; BD3-LM/MDLM 524B). Timing protocol as in the paragraph above. Batch 256 for ELF-B and GPT-2 small comes from a separate batch-scaling run (ELF-B at batch 32 in that run: 46.7 documents per second).

	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 $K{=}64$	28.5	34.5	39.2	48.4	35.0	32.3	38.1
PSF $K{=}32$	23.4	28.3	32.2	40.6	28.0	25.7	26.7 (extr.)
PSF $K{=}16$	17.7	21.8	24.9	31.4	22.0	20.2	22.2
PSF $K{=}8$	15.9	19.7	22.8	27.64 (5.237)	20.1	18.4	20.4
states-only, $p{=}0.6$	—	—	19.4	—	16.7	15.1	17.1

Appendix

Table 15: Generative perplexity under three judges, at four entropies and on a 512-token window. GPT-2 Large and Qwen columns: log-linear interpolation in each judge’s own tokens; window: first 512 GPT-2 Large tokens, matched at ELF-B’s 32-step window entropy 4.845 (Appendix D.2 ). Values in parentheses are measured at the entropy given; “extr.” = window bracket extrapolated by at most 0.06 nats; “—” = not computed or not reached without extrapolation;

p

= nucleus threshold; indented rows use the states-only prior of the

K{=}8

model above them; other rows of ours use the plain sampler.

System	Distinct-4 $\uparrow$	Self-BLEU-4 $\downarrow$	13-gram overlap
Held-out OpenWebText	0.905	0.148	2.67%
PSF $K{=}8$ , $t{=}0.75$	0.803	0.301	0.54%
PSF-L $K{=}8$ , $t{=}0.70$	0.781	0.296	0.83%
PSF $K{=}8$ , states-only, nucleus 0.55, ban	0.750	0.404	0.49%
PSF-L $K{=}8$ , states-only, nucleus 0.6	0.741	0.363	0.96%
ELF-B, 32 steps	0.801	0.292	0.04%

Appendix

Table 16: Cross-sample diversity and 13-gram overlap with a 2M-document training index. Self-BLEU uses 300 documents; overlap is the fraction of each dump’s T5-token 13-grams found in the index.

Sampler	$n$	Entropy	PPL with ban	PPL without ban	$\Delta$ PPL (%)	Loop, no ban
PSF $K{=}8$ , prior’s codes
$t{=}0.65$	4	5.111	20.25	20.2	$0$	9%
$t{=}0.70$	4	5.162	22.76	22.9	$-0.7$	6.9%
$t{=}0.70$	3	5.235	26.25	27.5	$-4.5$	—
$t{=}0.70$	6	5.120	20.50	20.6	$-0.5$	—
$t{=}0.70$	8	5.109	19.82	20.1	$-1.4$	—

Appendix

Table 17: Document-wide

n

-gram ban at sampling time. “PPL without ban” is the unbanned sweep interpolated at the banned run’s entropy; “Loop, no ban” is the unbanned model’s loop rate at the nearest setting (“—” = not measured). Every banned run has a loop rate of 0%.

Figure 11: Lexical topic-persistence measures.

J_{1}

compares adjacent-chunk word overlap with cross-document overlap; switch rate counts low-overlap transitions. All our rows use

K=8

p

: nucleus threshold; states: states-only prior; ban: 4-gram ban;

b

: lexical bonus. Dashed lines mark held-out text. Right: GPT-2 Large perplexity and entropy

H

; error bar /

\pm

: three-seed standard deviations.

	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

Appendix

Table 18: Blinded pairwise LLM judgement, 200 document pairs per row, each shown in both slot orders. Cells: consistent pairs won by system 1 : won by system 2; the remaining pairs disagree between orders. The small “states-only + ban” configuration is the states-only PSF row of Table 1 ; the PSF-L configuration here also uses the ban, unlike its main-table row.

Figure 12: Prompted continuation of 500 held-out documents from 256-token prompts. Relevance is content-word overlap above chance. Our filled points use inpainted plans; hollow points use oracle plans; all use

K=8

r

: resampling passes;

b

: lexical bonus. Dashed: real continuation. Right: continuation perplexity and entropy

H

under GPT-2 Large.

Text Generation

Text generation studies methods for producing coherent, useful language from prompts, data, and dialogue context.

See all 423 papers

Denoising Blocks, Not Tokens: Efficient Compressed Continuous Diffusion with Branching Token Realization
Oct 7, 2026Xinsong Feng, Peng Du, Zhizhuo Yang +3Diffusion Language Models Diffusion Language Model Inference
Georgia Institute of Technology · Writer AI Research · William & Mary
Masked Diffusion Decoding as $x$ -Prediction Flow
Jun 27, 2026Weitian Wang, Lianlei Shan, Shubham Rai +2Masked Diffusion Models Diffusion Language Model Decoding
Robert Bosch GmbH, Germany · University of the Chinese Academy of Sciences, China · Ruhr University Bochum, Germany
How to Train Your Latent Diffusion Language Model Jointly With the Latent Space
May 8, 2026Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov +4Non-Autoregressive Text Generation Latent Diffusion Models
Constructor University, Bremen, Germany · HSE University, Moscow, Russia · Applied AI Institute, Moscow, Russia +2

Oct 7, 2026J/K move · Enter open · S save

Denoising Blocks, Not Tokens: Efficient Compressed Continuous Diffusion with Branching Token Realization

Xinsong Feng, Peng Du, Zhizhuo Yang, Daniel M. Bikel+2

Georgia Institute of Technology · Writer AI Research · William & Mary

Diffusion Language Models Diffusion Language Model Inference Diffusion Model Acceleration Continuous Diffusion Language Models Diffusion Model Compression

Open paper AskPDF arXiv

Diffusion language models (DLMs) generate text through iterative parallel refinement, offering the potential for higher throughput than autoregressive (AR) decoding. However, most DLMs still maintain one generative state per token, so every denoising step processes a state sequence as long as the output sequence, limiting the throughput gains from parallel generation. Continuous DLMs provide an additional degree of freedom: a single continuous state can represent multiple tokens, allowing diffusion to operate on a much shorter latent sequence. We introduce Branching Latent Diffusion (BLD), which exploits this flexibility by compressing a 1024-token sequence into only 64 block latents, a

16\times

reduction. BLD combines latent compression with branching token realization, where each latent is decoded by a local AR branch and all branches run in parallel. Because strong compression makes joint latent generation difficult, BLD generates the latents in groups, conditioning each group on previously generated latents. In end-to-end evaluation on the same GPU, BLD reduces generation FLOPs by more than

80\times

and increases throughput by more than

6\times

relative to the similarly sized ELF-L baseline. Compared with the AR baseline, BLD achieves more than

6\times

higher throughput and more than

4\times

lower latency. Despite the compression, BLD maintains competitive local fluency and diversity, although long-range coherence remains challenging. Overall, BLD shows that moving diffusion from token-level states to compressed latent sequences can substantially improve the efficiency of long-sequence generation.

Abstract

Text Generation

Explore similar work

Denoising Blocks, Not Tokens: Efficient Compressed Continuous Diffusion with Branching Token Realization

Masked Diffusion Decoding as xxx-Prediction Flow

How to Train Your Latent Diffusion Language Model Jointly With the Latent Space

Masked Diffusion Decoding as $x$ -Prediction Flow