---
license: mit
library_name: pytorch
language:
  - en
datasets:
  - embedded-language-flows/openwebtext-t5
pipeline_tag: text-generation
tags:
  - diffusion-language-model
  - flow-matching
  - parallel-decoding
  - arxiv:2610.12898
---

# Predictive-State Flows (PSF): checkpoints

Checkpoints for **Rethinking the Roles of Diffusion and Decoding: Predictive-State Flows for Parallel Language
Generation** ([paper](https://arxiv.org/abs/2610.12898),
[code](https://github.com/jinjinhe2001/Predictive-State-Flows-for-Parallel-Language-Generation)).

A predictive state s<sub>k</sub> summarizes only the blocks before block k. A rectified-flow prior generates the K
predictive states of a 1024-token document jointly from noise; a block decoder then writes all K blocks in parallel,
autoregressively only within each block. Unconditional generation, trained on OpenWebText with the T5 tokenizer.

## Files

| File | Model | K | Gen. PPL @ entropy | MAUVE@512 | TFLOPs / doc |
|---|---|---|---|---|---|
| `psf_k64.pt` | PSF | 64 | 41.15 @ 5.183 | 0.87 | 0.80 |
| `psf_k32.pt` | PSF | 32 | 33.56 @ 5.177 | 0.89 | 0.52 |
| `psf_k16.pt` | PSF | 16 | 25.67 @ 5.173 | 0.90 | 0.38 |
| `psf_k8.pt` | PSF | 8 | 22.75 @ 5.164 | 0.93 | 0.31 |
| `psf_k8_states-only.pt` | PSF, states-only prior | 8 | 18.95 @ 5.158 | 0.85 | 0.31 |
| `psf-l_k64.pt` | PSF-L | 64 | 42.2 @ 5.16 (interpolated) | | |
| `psf-l_k32.pt` | PSF-L | 32 | 27.1 @ 5.16 (interpolated) | | |
| `psf-l_k16.pt` | PSF-L | 16 | 21.22 @ 5.181 | 0.91 | 1.35 |
| `psf-l_k8.pt` | PSF-L | 8 | 17.79 @ 5.146 | 0.90 | 1.10 |
| `psf-l_k8_states-only.pt` | PSF-L, states-only prior | 8 | 15.98 @ 5.144 | 0.89 | 1.10 |
| `ar-197m.pt` | same-data autoregressive baseline | - | 15.37 @ 5.151 | 0.95 | 0.43 |

PSF: 197M generation parameters (block decoder 109M + prior 88M), 0.93 GB per file. PSF-L: 677M (369M + 308M),
3.15 GB per file. Both files also carry the encoder. AR: 197M, 0.79 GB. All fp32 EMA weights. Numbers are from
Table 1 of the paper (K = 64 and 32 of PSF-L from its block-length table): 1000 samples, GPT-2 Large judge, with the
sampling settings in `configs/<name>.sampling.args`.

## Usage

```bash
git clone https://github.com/jinjinhe2001/Predictive-State-Flows-for-Parallel-Language-Generation.git
cd Predictive-State-Flows-for-Parallel-Language-Generation && pip install -r requirements.txt
hf download jinjinhe2001/PSF-ParallelLM --include "psf_k8.pt" --local-dir checkpoints    # or omit --include for all
python -m scripts.eval_gen --ckpt checkpoints/psf_k8.pt $(cat configs/models/psf_k8.args) \
    $(cat configs/sampling/psf_k8.args) --save-texts outputs/texts_psf_k8
```

The architecture is not stored in the PSF checkpoints: `configs/<name>.model.args` (also in the code repository)
holds the flags that rebuild it. Each PSF file is `{"model": state_dict, "steps": int}`; the AR file stores its own
architecture as `{"model": state_dict, "ar_config": dict, "step": int}`. All files load with
`torch.load(path, weights_only=True)`. The AR baseline is sampled with `scripts.sample_ar` (see the code README).

## Citation

```bibtex
@article{he2026psf,
  title         = {Rethinking the Roles of Diffusion and Decoding: Predictive-State Flows for Parallel Language Generation},
  author        = {He, Jinjin and Feng, Xinsong and Fu, Lucheng and Wang, Sinan and Zhu, Bo and Wang, Jiayun},
  journal       = {arXiv preprint arXiv:2610.12898},
  year          = {2026}
}
```