README.md

<h2 align="center">
Beyond Token-Level Cross-Entropy: Fréchet Distributional Post-Training for Autoregressive Image Generation
</h2>

<p align="center">
  <a href="https://scholar.google.com/citations?user=tyYxiXoAAAAJ">Jinhua Zhang</a><sup>*</sup>,
  <a href="https://openreview.net/profile?id=~Yisong_Lin1">Yisong Lin</a><sup>*</sup>,
  <a href="https://scholar.google.com/citations?user=CsVTBJoAAAAJ">Wei Long</a>,
  <a href="https://scholar.google.com/citations?user=-kSTt40AAAAJ">Shuhang Gu</a><sup>†</sup>
</p>

<p align="center">
  University of Electronic Science and Technology of China
</p>

<p align="center">
  <sup>*</sup> Equal contribution &nbsp;&nbsp;
  <sup>†</sup> Corresponding author
</p>

<p align="center">
  <a href="https://arxiv.org/abs/2608.00562">
    <img src="https://img.shields.io/badge/arXiv-2608.00562-b31b1b.svg" alt="arXiv">
  </a>
  <a href="https://github.com/CVL-UESTC/FDPT-AR">
    <img src="https://img.shields.io/github/stars/CVL-UESTC/FDPT-AR?style=social" alt="GitHub stars">
  </a>
  <a href="https://huggingface.co/CVLUESTC/FDPT-AR">
    <img src="https://img.shields.io/badge/🤗%20Hugging%20Face-CVLUESTC%2FFDPT--AR-yellow" alt="Hugging Face">
  </a>
</p>

<p align="center">
  ⭐ If you find this work useful, please consider giving this repository a star!
</p>

---

## ✨ Key Contributions

- **Image-level post-training.** We optimize pretrained discrete autoregressive
  generators using an EMA-based Fréchet distributional objective without
  retaining cross-entropy.

- **Detached rollout context replay.** Model-generated rollouts provide
  inference-aligned replay contexts, reducing the mismatch introduced by
  teacher forcing.

- **Differentiable discrete decoding.** A probability-level straight-through
  estimator preserves hard argmax decoding in the forward pass while
  propagating image-level gradients through soft probabilities.

- **Plug-and-play improvement.** FDPT-AR improves LlamaGen, TiTok, GigaTok,
  and VAR without changing their architectures, parameter counts, or inference
  procedures.

<p align="center">
  <img src="asset/ema_fd_framework.png" width="90%" alt="FDPT-AR framework">
</p>

---

## ⚙️ Environment

We recommend using Python 3.10 and NVIDIA GPUs.

```bash
conda create -n fdpt-ar python=3.10 -y
conda activate fdpt-ar
pip install -r requirements.txt
```

Download all pretrained checkpoints and evaluation statistics:

```bash
bash scripts/download_assets.sh all
```

To download only one model family:

```bash
bash scripts/download_assets.sh llamagen
bash scripts/download_assets.sh stats
```

---

## 🤗 Post-trained Checkpoints

Our FDPT-AR post-trained generator weights are available on
[Hugging Face](https://huggingface.co/CVLUESTC/FDPT-AR).

| Model | Post-trained checkpoint |
| --- | --- |
| LlamaGen-B | [`llamagen-b.pt`](https://huggingface.co/CVLUESTC/FDPT-AR/blob/main/llamagen-b.pt) |
| TiTok-L-32 | [`titok-l32.pt`](https://huggingface.co/CVLUESTC/FDPT-AR/blob/main/titok-l32.pt) |
| TiTok-B-64 | [`titok-b64.pt`](https://huggingface.co/CVLUESTC/FDPT-AR/blob/main/titok-b64.pt) |
| GigaTok-S-S | [`gigatok-ss.pt`](https://huggingface.co/CVLUESTC/FDPT-AR/blob/main/gigatok-ss.pt) |
| VAR-d16 | [`var-d16.pt`](https://huggingface.co/CVLUESTC/FDPT-AR/blob/main/var-d16.pt) |
| VAR-d20 | [`var-d20.pt`](https://huggingface.co/CVLUESTC/FDPT-AR/blob/main/var-d20.pt) |
| VAR-d24 | [`var-d24.pt`](https://huggingface.co/CVLUESTC/FDPT-AR/blob/main/var-d24.pt) |

Download all released post-trained checkpoints:

```bash
bash scripts/download_assets.sh released
```

Download a single checkpoint:

```bash
hf download CVLUESTC/FDPT-AR \
  llamagen-b.pt \
  --local-dir checkpoints/fdpt-ar
```

These files contain the FDPT-AR post-trained generator weights. The corresponding
pretrained tokenizer/VAE checkpoints and evaluation statistics can be downloaded
using `scripts/download_assets.sh`.

---

## 🔥 Training

```bash
bash scripts/train.sh \
  --model llamagen-b \
  --ckpt-dir checkpoints/llamagen \
  --bs 8
```

The script automatically generates the initialization images and saves the
post-training checkpoints to:

```text
outputs/train/<model>/
```

Available models and their required checkpoint files are listed below.

| Model | `--model` | Files in `--ckpt-dir` |
| --- | --- | --- |
| LlamaGen-B | `llamagen-b` | `c2i_B_256.pt`, `vq_ds16_c2i.pt` |
| LlamaGen-L | `llamagen-l` | `c2i_L_256.pt`, `vq_ds16_c2i.pt` |
| TiTok-L-32 | `titok-l32` | `generator_titok_l32.bin`, `tokenizer_titok_l32.bin` |
| TiTok-B-64 | `titok-b64` | `generator_titok_b64.bin`, `tokenizer_titok_b64.bin` |
| VAR-d16 | `var-d16` | `var_d16.pth`, `vae_ch160v4096z32.pth` |
| VAR-d20 | `var-d20` | `var_d20.pth`, `vae_ch160v4096z32.pth` |
| VAR-d24 | `var-d24` | `var_d24.pth`, `vae_ch160v4096z32.pth` |
| GigaTok-S-S | `gigatok-ss` | `GPT_B256_e300_VQ_SS.pt`, `VQ_SS256_e100.pt` |

To select specific GPUs, set `CUDA_VISIBLE_DEVICES` before running the command:

```bash
CUDA_VISIBLE_DEVICES=0,1 bash scripts/train.sh \
  --model llamagen-b \
  --ckpt-dir checkpoints/llamagen \
  --bs 8
```

---

## 🖼️ Sampling

Generate 50,000 images using a locally trained or released post-trained checkpoint:

```bash
bash scripts/sample.sh \
  --model llamagen-b \
  --ckpt-dir checkpoints/llamagen \
  --bs 8
```

Generated images are saved to:

```text
outputs/samples/<model>/
```

The script uses the latest checkpoint in `outputs/train/<model>/`. If no local
checkpoint is found, it automatically loads
`checkpoints/fdpt-ar/<model>.pt`.

---

## 📊 Evaluation

Compute FID and FDr6 for the generated images:

```bash
bash scripts/evaluate.sh \
  --model llamagen-b \
  --bs 16
```

The evaluation results are saved to:

```text
outputs/eval/<model>.json
```

---

## 🤝 Acknowledgements

This repository is built upon the following excellent projects:

- [LlamaGen](https://github.com/FoundationVision/LlamaGen)
- [TiTok/1D-Tokenizer](https://github.com/bytedance/1d-tokenizer)
- [VAR](https://github.com/FoundationVision/VAR)
- [GigaTok](https://github.com/SilentView/GigaTok)
- [FD-Loss](https://github.com/Jiawei-Yang/FD-Loss)

We sincerely thank the authors for their outstanding work.

---

## 📝 Citation

If you find this work useful, please consider citing:

```bibtex
@misc{zhang2026tokenlevelcrossentropyfrechetdistributional,
  title={Beyond Token-Level Cross-Entropy: Fr\'echet Distributional Post-Training for Autoregressive Image Generation},
  author={Jinhua Zhang and Yisong Lin and Wei Long and Shuhang Gu},
  year={2026},
  eprint={2608.00562},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2608.00562}
}
```

---

## 📄 License

FDPT-AR-specific modifications are released under the Apache-2.0 License.
Upstream-derived files retain their original licenses. See
[`THIRD_PARTY_NOTICES.md`](THIRD_PARTY_NOTICES.md) for details.

---

## 📩 Contact

For questions or collaborations, please contact
[Jinhua Zhang](mailto:jinhua.zjh@gmail.com) or
[Yisong Lin](mailto:yisongl164@gmail.com).