Beyond Token-Level Cross-Entropy: Fréchet Distributional Post-Training for Autoregressive Image Generation
Jinhua Zhang*,
Yisong Lin*,
Wei Long,
Shuhang Gu†
University of Electronic Science and Technology of China
* Equal contribution
† Corresponding author
⭐ If you find this work useful, please consider giving this repository a star!
---
## ✨ 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.
---
## ⚙️ 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//
```
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//
```
The script uses the latest checkpoint in `outputs/train//`. If no local
checkpoint is found, it automatically loads
`checkpoints/fdpt-ar/.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/.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).