pipeline_tag: unconditional-image-generation
Parallel Rollout Approximation (PRA)
This repository contains the weights for Parallel Rollout Approximation (PRA), a scalable framework for class-conditional pixel-space autoregressive image generation.
More details can be found in the paper Parallel Rollout Approximation for Pixel-Space Autoregressive Image Generation.
- Repository: GitHub Repository
- Paper: arXiv:2606.27978
Model Description
Parallel Rollout Approximation (PRA) is a pixel-space continuous-token autoregressive (AR) generation model. PRA generates low-dimensional intermediate states instead of high-dimensional pixel patches, mapping them back to pixel-space tokens with a pixel decoder. It effectively mitigates error accumulation during autoregressive steps by approximating the pixel-feedback interface encountered during inference-time rollout while retaining parallel teacher-forced training.
Model Checkpoints
The following checkpoints are available:
| Model | Params | FID (256x256) | Weight |
|---|---|---|---|
| PRA-S | 135M | 2.58 | PRA_S.pt |
| PRA-B | 250M | 2.21 | PRA_B.pt |
| PRA-L | 511M | 1.94 | PRA_L.pt |
Environment & Usage
For environment setup, training, and evaluation scripts, please refer to the official GitHub Repository.
Sampling Example
You can run distributed class-balanced sampling using the sample_ddp.py script provided in the repository:
ckpt=your_ckpt_path
sample_dir=your_result_path
torchrun --nnodes=1 --nproc_per_node=4 --node_rank=0 \
sample_ddp.py \
--ckpt $ckpt \
--sample-dir $sample_dir \
--model PRA-L \
--image-size 256 \
--patch-size 16 \
--latent-dim 16 \
--cfg-scale 4.1 \
--sample-steps 100 \
--sampler euler_maruyama \
--per-proc-batch-size 200 \
--sample-mask-rate 0.9 \
--token-mask-rate 0.5 \
--save-png
Citation
@article{xu2026parallel,
title={Parallel Rollout Approximation for Pixel-Space Autoregressive Image Generation},
author={Xu, Jiayi and He, Di and Ke, Guolin},
journal={arXiv preprint arXiv:2606.27978},
year={2026}
}