ICLR 2026 Paper
Pheng-Ann Heng3
Stefano Ermon2
James Zhou2
Angela Yao1
Overview: We propose reAR, a simple regularization method that fixes generator-tokenizer inconsistency in visual autoregressive models, i.e, let the AR model generates tokens that is more friendly for the decoder of visual tokenizer. Without altering tokenization and inference, reAR boosts AR on vanilla VQGAN to 1.86 (461M) and AliTok to 1.42 (177M).
# Create conda environment
conda create -n rear python=3.10
conda activate rear
# Install dependencies
pip install -r requirements.txtSet up your project directory as follows:
re-ar/
├── configs/ # Configuration files
├── dataset/ # Dataset directory
│ ├── pretokenized/ # Pre-tokenized data (recommended)
│ │ └── maskgitvq.jsonl # Download from HuggingFace
│ └── imagenet_shard/ # Original ImageNet (optional)
│ ├── train/
│ │ ├── imagenet-train-0000.tar
│ │ ├── imagenet-train-0001.tar
│ │ └── ...
│ └── val/
│ ├── imagenet-val-0000.tar
│ ├── imagenet-val-0001.tar
│ └── ...
├── ckpt/ # Model checkpoints
│ └── maskgitvq.bin # Download from HuggingFace
└── scripts_bash/ # Training and evaluation scripts
Download the pre-tokenized files for faster training:
# Download pre-tokenized data
wget https://huggingface.co/yucornetto/RAR/resolve/main/maskgitvq.jsonl -O dataset/pretokenized/maskgitvq.jsonlFor original ImageNet in webdataset format, follow the TiTok instructions.
Download the MaskGiT-VQGAN tokenizer:
# Download tokenizer checkpoint
wget https://huggingface.co/fun-research/TiTok/resolve/main/maskgit-vqgan-imagenet-f16-256.bin -O ckpt/maskgitvq.binAlternatively, for direct evaluation, the following checkpoints are available:
| Model | Checkpoints |
|---|---|
| reAR-S | reAR-S.safetensors |
| reAR-B | reAR-B.safetensors |
| reAR-L | reAR-L.safetensors |
bash scripts_bash/rear_train.shIf you prefer to run training manually, here's the complete setup:
# Navigate to project directory
cd ~/re-ar
# Activate environment
source ~/.bashrc
eval "$(conda shell.bash hook)"
conda activate rear
# Set environment variables
export PYTHONPATH=$(pwd)
export WANDB_INIT_TIMEOUT=300
# Configure training parameters
entity="your_wandb_entity" # Replace with your Weights & Biases entity
config_name='rear_l'
# Launch training
accelerate launch \
--num_machines=1 --num_processes=8 --machine_rank=0 \
--main_process_ip=127.0.0.1 --main_process_port=9999 --same_network \
scripts/train_rear.py config="configs/${config_name}.yaml" \
experiment.entity="${entity}" \
experiment.output_dir="temp/${config_name}"# Clone ADM evaluation repository
git clone https://github.com/openai/guided-diffusion.git
# Download reference batches
wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/VIRTUAL_imagenet256_labeled.npz# Quick evaluation
bash scripts_bash/rear_test.sh# Navigate to project directory
cd ~/re-ar
# Activate environment
source ~/.bashrc
eval "$(conda shell.bash hook)"
conda activate rear
# Set environment variables
export PYTHONPATH=$(pwd)
export WANDB_INIT_TIMEOUT=300
# Configure evaluation parameters
config_name="rear_l" # Change to your config
output_dir="${config_name}"
checkpoint_path="path/to/your/checkpoint" # Specify your checkpoint path
# Generate samples
torchrun --nnodes=1 --nproc_per_node=8 --rdzv-endpoint=localhost:19999 \
scripts/sample_imagenet_rear.py config="configs/training/generator/${config_name}.yaml" \
experiment.output_dir="${output_dir}" \
experiment.generator_checkpoint="${checkpoint_path}"
# Evaluate samples
python3 guided-diffusion/evaluations/evaluator.py \
VIRTUAL_imagenet256_labeled.npz ${output_dir}.npzWe thank the following projects for their excellent work:
- RAR: Randomized Autoregressive Visual Generation
- TiTok: An Image is Worth 32 Tokens for Reconstruction and Generation
- MaskBit: Embedding-free Image Generation via Bit Tokens
- REPA: Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think
- AliTok: Towards Sequence Modeling Alignment between Tokenizer and Autoregressive Model