You can deploy this model and start earning money today!

CVLUESTC/
FDPT-AR

2608.00562

Views

No views yet

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

arXiv

⭐ 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 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.

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:

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:

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:

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:

outputs/eval/<model>.json

🤝 Acknowledgements

This repository is built upon the following excellent projects:

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

📩 Contact

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