jinjinhe2001/PSF-ParallelLM

Predictive-State Flows (PSF): checkpoints

Checkpoints for Rethinking the Roles of Diffusion and Decoding: Predictive-State Flows for Parallel Language Generation (paper, code).

A predictive state s_k 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

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

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