Should diffusion priors generate predictive states instead of token-level latents in language models?

The prior and decoder should split the work

The paper asks which information a continuous diffusion prior should be responsible for, and which decisions should be left to the decoder. In many existing continuous diffusion language models, the latents produced by the prior are mapped almost deterministically to tokens, either by a decoder or through a token embedding matrix. That arrangement puts nearly all generative work on the prior: it must specify fine-grained discrete token choices inside a continuous space, while the decoder mostly recovers choices already encoded in those states.

The authors argue this is a poor fit. In image generation, perceptual compression can discard visually negligible detail and represent the remaining variation continuously. Reconstruction-based text latents, by contrast, tend to preserve the fine-grained discrete distinctions that separate one plausible continuation from another. The prior must then model those distinctions and their dependencies jointly in continuous space, which the paper describes as making the joint latent distribution difficult to learn.

The proposed division is therefore: the diffusion prior generates predictive information that guides what comes next, and the decoder resolves the remaining local token choices. The rest of the paper tests whether that split actually helps.

Predictive states see only the preceding context

Predictive-State Flows (PSF) partitions a sequence of length L into K consecutive blocks. Each block k is associated with a continuous state s_k computed solely from the tokens before that block, so the state summarizes the preceding context without observing the block it will help generate. This is the key contrast with a content latent, which encodes the current block itself and therefore carries both predictive information and the specific token choices inside it.

A single block-causal encoder pass produces pooled states, and shifting them by one block yields the state sequence, with the first state set to a beginning-of-sequence state. Because each state is built from a prefix, the prior never has to commit to the exact tokens of the block it precedes.

The paper also trains an optional realization-code channel: a second encoder head defines a posterior over a 64-dimensional code per block, regularized toward a standard normal with a budget of one nat per token. These codes give the decoder extra local information, and the paper later tests what happens when they are removed.

A rectified-flow prior over K states, not L tokens

A rectified-flow prior generates the trajectory of predictive states jointly. It operates on a short sequence of K states with full self-attention across them, transforming Gaussian noise into the full state trajectory in a fixed number of sampling steps; the main configuration uses 50 Euler steps of the probability-flow ODE. The prior is trained with an x-prediction velocity loss and a logit-normal time distribution.

The compression is the point. Diffusion is reduced from L token-level states to K predictive states, with K much smaller than L, and the token-level sequential depth falls to L/K. For a 1024-token document with K=8, the prior handles 8 states rather than 1024 positions.

A block-causal decoder then generates all K blocks in parallel, autoregressively only within each block. Each branch reads the cumulative memory of states up to its own block and advances synchronously for L/K steps. The paper reports a serial depth of 50 + L/K, compared with L steps for standard autoregressive generation, where the 50 counts the prior's sampling steps.

One-token landing targets keep parallel blocks coupled

Generating every block independently from a shared plan would leave the blocks unconnected. PSF adds a lightweight coupling: after the first decoding step, each branch conditions on the first token of the following block, called a landing target. The first decoding step produces these landing targets, and from step two onward each branch continues with knowledge of how the next block actually begins.

Training has to match this inference behavior. The paper uses self-realized prefixes, replacing states per block with probability 0.5 by states computed from the model's own rollouts at temperature 0.65 and nucleus 0.95, so the decoder learns to work with states that came from generation rather than from clean text.

The authors are explicit that this coupling is thin. Blocks are decoded in parallel from a shared plan and connected only by one-token landing targets, so whether many blocks stay consistent in topic, entities and discourse over longer documents is untested. They note that topic persistence is already the weakest aspect of their samples and list stronger coupling as a next step.

Fewer FLOPs at lower generative perplexity

The headline comparison is unconditional generation on OpenWebText. A 197M-parameter PSF model with K=8 reaches generative perplexity 22.8 at entropy 5.16, using 0.31 TFLOPs per 1024-token document. ELF-B with 32 sampling steps reaches 24.0 at 7.2 TFLOPs. Generative perplexity here is the exponential of the corpus-level mean token negative log-likelihood under a judge model, and entropy is the unigram entropy of the judge's tokens averaged per document, so the two numbers must be read together: the comparison is at matched entropy, not at arbitrary sampling settings.

At those measured operating points, MAUVE@512 is 0.93 for PSF and 0.73 for ELF-B. MAUVE measures distributional closeness to reference text over a 512-token window, a different question from perplexity, and the paper reports both rather than substituting one for the other.

The discussion states that with only K=8 predictive states for a 1024-token document, PSF uses 23 times fewer FLOPs than 32-step ELF-B while achieving lower generative perplexity, higher MAUVE and higher batched throughput. The paper also notes that the FLOP saving becomes a throughput advantage only at large batch, so the efficiency claim is conditional on batch size.

Predictive states beat states that encode their own block

The most direct test of the paper's thesis is a controlled comparison between predictive states and content states. Content states additionally encode their own blocks, so they carry the token choices the prior is supposed to leave to the decoder. Across the controlled comparisons, content states yield 8 to 16 percent higher generative perplexity than predictive states.

A second control removes the sampling of state trajectories: replacing sampled predictive-state trajectories with a shared mean lowers MAUVE@512 by 0.11. This isolates the value of generating a varied trajectory rather than a single averaged plan, and it is measured on MAUVE rather than perplexity.

The paper also reports that increasing block length improves PSF even as it degrades block diffusion. Together these comparisons support the authors' interpretation that what matters is not only how compactly text is represented but how generative roles are divided between prior and decoder.

Dropping realization codes trades coherence for perplexity

For K=8 the paper trains a states-only prior that generates only predictive states, using zero code targets while keeping the code-trained block decoder and fixing its code inputs to zero. When the prior generates predictive states alone, nucleus sampling reduces matched-entropy perplexity to roughly 20 percent below ELF-B.

That gain is not free. The paper states that the states-only configuration trades topic persistence, diversity and judged coherence for lower perplexity. A reader comparing configurations therefore cannot treat the perplexity improvement as a uniform win; it depends on which output property is being optimized.

The training recipe itself is multi-stage, with separate stages for the base model, landing targets, self-realized prefixes, realization codes and each block length, and each prior is warm-started from the previous block length. The authors acknowledge that the benefit of realization codes is only partly understood and that a simpler recipe remains future work.

The evidence covers small models and unconditional text

All experiments use at most 677M generation parameters, a fixed 1024-token canvas and unconditional OpenWebText generation. Larger models and longer sequences are unexplored, so the reported FLOP and quality comparisons should not be assumed to hold at other scales or sequence lengths.

Prompted continuation works without retraining but uses the prompt weakly, which the authors say calls for a prompt-conditioned prior. Autoregressive models remain ahead, although the gap to the same-data baseline narrows at the largest scale tested. Generative perplexity also misses long-range coherence, which the paper lists as something still to be measured.

Given those boundaries, the paper's most decision-relevant claim is a comparison rather than a universal recipe: at matched entropy on 1024-token unconditional OpenWebText, K=8 predictive states delivered lower generative perplexity and higher MAUVE@512 than 32-step ELF-B at a fraction of the FLOPs, and predictive states beat content states by 8 to 16 percent generative perplexity. The priority the evidence supports is designing the latent around which information the prior can generate easily and the decoder can resolve, not around how completely the latent reconstructs the text.

Design latents for the role split, not reconstruction

The paper's closing argument is that a useful latent need not preserve every detail of the output. It can instead represent information that is easy for a continuous prior to generate and sufficient for a decoder to turn into coherent text. Predictive states are one realization of that principle, and the controlled comparisons are what make the argument more than a design preference.

For readers weighing continuous diffusion language models, the practical split the evidence supports is this: if the goal is cheap parallel generation at matched entropy on short unconditional documents, generating K predictive states with a flow prior and letting a block-causal decoder resolve tokens is the configuration the paper measured favorably. If the goal is topic persistence, diversity or judged coherence, the states-only variant's trade-off and the weak one-token coupling between blocks are the conditions to check first.

Sources used in this answer