---
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](https://huggingface.co/papers/2606.27978).

* **Repository:** [GitHub Repository](https://github.com/MangataX/PRA)
* **Paper:** [arXiv:2606.27978](https://huggingface.co/papers/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](https://huggingface.co/MangataX/PRA/blob/main/PRA_S.pt) |
| PRA-B | 250M | 2.21 | [PRA_B.pt](https://huggingface.co/MangataX/PRA/blob/main/PRA_B.pt) |
| PRA-L | 511M | 1.94 | [PRA_L.pt](https://huggingface.co/MangataX/PRA/blob/main/PRA_L.pt) |

## Environment & Usage

For environment setup, training, and evaluation scripts, please refer to the official [GitHub Repository](https://github.com/MangataX/PRA).

### Sampling Example

You can run distributed class-balanced sampling using the `sample_ddp.py` script provided in the repository:

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

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