| --- |
| 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} |
| } |
| ``` |
|
|
