CVLUESTC/FDPT-AR

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.

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 scripts/download_assets.sh all

To download only one model family:

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.

Model	Post-trained checkpoint
LlamaGen-B	llamagen-b.pt
TiTok-L-32	titok-l32.pt
TiTok-B-64	titok-b64.pt
GigaTok-S-S	gigatok-ss.pt
VAR-d16	var-d16.pt
VAR-d20	var-d20.pt
VAR-d24	var-d24.pt

Download all released post-trained checkpoints:

bash scripts/download_assets.sh released

Download a single checkpoint:

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 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:

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:

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 scripts/sample.sh \
  --model llamagen-b \
  --ckpt-dir checkpoints/llamagen \
  --bs 8

Generated images are saved to:

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 scripts/evaluate.sh \
  --model llamagen-b \
  --bs 16

The evaluation results are saved to:

outputs/eval/<model>.json

---

🤝 Acknowledgements

This repository is built upon the following excellent projects:

• LlamaGen
• TiTok/1D-Tokenizer
• VAR
• GigaTok
• FD-Loss

We sincerely thank the authors for their outstanding work.

---

📝 Citation

If you find this work useful, please consider citing:

@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 for details.

---

📩 Contact

For questions or collaborations, please contact Jinhua Zhang or Yisong Lin.