wAR-Tok: Learning Discrete Autoregressive Priors with Wasserstein Gradient Flow

CIFAR-10 implementation of our paper.

Features

Discrete image tokenization with Wasserstein gradient flow regularization.
Transformer encoder and decoder with cross-attention, RMSNorm and GEGLU.
Six variants: baseline, tail dropout, flow matching, FM + tail dropout, FM + tied tail dropout, and wAR-Tok.
Two-stage training: tokenizer training followed by an autoregressive prior.
Mixed precision, multi-GPU training and optional W&B logging.

Quick Start

Installation

Use Python 3.12 and CUDA. The training environment uses PyTorch 2.8.0 / CUDA 12.8.

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Data Preparation

python scripts/prepare_cifar10.py --download

This prepares the CIFAR-10 training split under data/cifar10. Download the FID reference to refs/cifar10-32x32_local.npz.

Training

The default configuration uses 8 GPUs, global batch size 512 and BF16. Each method has a *_tokenizer.yaml and a *_ar.yaml in configs/cifar10: train the tokenizer for 400K steps, then the AR prior for 200K steps.

Stage 1: Train the tokenizer

python train.py --config=configs/cifar10/war_tokenizer.yaml

Stage 2: Train the AR prior

python train.py --config=configs/cifar10/war_ar.yaml \
  ae_ckpt_path=./outputs/war_tokenizer/000/ckpts/last.ckpt

The AR configuration defaults to the matching tokenizer's 000/ckpts/last.ckpt. Set WANDB_MODE=online to enable experiment logging.

Evaluation

Use the matching tokenizer configuration for reconstruction and the AR configuration for generation. FM reconstruction uses 25 Euler steps; generation uses a single decoder step. Evaluation uses FP32; generation uses AR EMA, temperature 1, CFG 0 and no top-k/top-p filtering.

Reconstruction

CUDA_VISIBLE_DEVICES=0 python evaluate.py \
  --config configs/cifar10/war_tokenizer.yaml --checkpoint checkpoints/war.pt \
  --mode reconstruct --data data/cifar10 --num-samples 50000 \
  --output outputs/war-reconstruction --fid-reference refs/cifar10-32x32_local.npz

Generation

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 torchrun --standalone --nproc-per-node=8 evaluate.py \
  --config configs/cifar10/war_ar.yaml --checkpoint checkpoints/war.pt \
  --num-samples 50000 --batch-size 64 --seed 42 --no-legacy \
  --output outputs/war-generation --fid-reference refs/cifar10-32x32_local.npz

Non-legacy mode (--no-legacy) uses a different seed on each GPU and generally yields lower gFID scores, without changing the relative ranking of the methods or the conclusions supported by the results.

To evaluate all six checkpoints with --no-legacy:

python scripts/evaluate_all.py --gpus 0,1,2,3,4,5,6,7 \
  --output outputs/cifar10-nonlegacy

Use --methods baseline war for a subset or --resume to reuse completed evaluations.

Exporting Checkpoints

python scripts/export_checkpoint.py --source path/to/raw-ar.ckpt \
  --config configs/cifar10/war_ar.yaml --output checkpoints/war.pt

The export includes tokenizer and raw/EMA AR weights without optimizer states.

Checkpoints

Name checkpoints after their configuration prefix (e.g. baseline_ar.yaml → checkpoints/baseline.pt, war_ar.yaml → checkpoints/war.pt); each includes the tokenizer and raw/EMA AR weights. Hugging Face repository: wARTok.

Method	rFID	Historical gFID	Non-legacy gFID	Download
Baseline	2.3015	52.5343	49.4191	Download
Tail Dropout	3.3567	23.3713	19.8997	Download
FM	3.4194	38.6309	35.8593	Download
FM + Tail	4.4247	22.7613	19.0569	Download
FM + Tail + Tied	6.4911	44.7971	41.1334	Download
wAR-Tok	3.0648	11.3059	7.3749	Download

Download the FID reference to refs/cifar10-32x32_local.npz.

Project Structure

wAR/
├── configs/cifar10/   # Training configurations
├── checkpoints/      # Downloaded weights
├── scripts/          # Data preparation and evaluation tools
├── models.py         # Model architectures
├── train.py          # Tokenizer and AR training
├── evaluate.py       # Reconstruction and generation
├── dataset.py        # Data loading
├── lpips.py          # Perceptual loss
└── calc_fid.py        # ADM FID evaluation

References

This project builds on EDM, ADM, and LPIPS.

License

Project-authored code is covered by MIT. Third-party components retain their original licenses.