XYZFlow

Official implementation of XYZFlow: Scaling Multidimensional Shortcut Flows for Efficient Generative Modeling.

High-fidelity image generation faces a trade-off between speed and quality. Diffusion models produce strong visuals but require costly iterative sampling. Existing efficient methods mainly distill pretrained models into few-step samplers, a challenging process that depends heavily on teacher-model quality. In this paper, we introduce XYZFlow, a framework that rethinks efficient generation through multidimensional scaling of flow matching. Unlike single-step mappings, XYZFlow enhances expressivity by making probability paths more identifiable and learnable through structured multidimensional conditioning. We view autoregressive modeling as implicit flow straightening, where richer context reduces trajectory ambiguity. XYZFlow realizes this idea through two orthogonal dimensions: temporal scaling, which uses non-Markovian conditioning on the full denoising history; and spatial scaling, enabled by Next Shortcut Prediction, which sequentially generates patches using preceding patches' denoising trajectories as priors. Experiments show that XYZFlow achieves state-of-the-art performance, with 7.2-8.5x teacher speedups and competitive FID, while Next Shortcut Prediction delivers superior quality-latency trade-offs over model scaling or step reduction.

Highlights

• Teacher trajectory generation from pretrained teacher checkpoints
• XYZFlow training with teacher-forcing and patch-scheduled supervision
• Patch-scheduled image generation
• Uniform-step baseline evaluation

News

• Paper: docs/XYZFlow_paper.pdf
• Project page: https://spherelab.ai/xyzflow

Installation

Install PyTorch for your CUDA version first, then install the remaining dependencies:

pip install -r requirements.txt

flash-attn may require a local CUDA toolkit compatible with the installed PyTorch version.

Checkpoints

xAR Base Models

XYZFlow uses pretrained xAR checkpoints as the teacher/backbone models. Download them from Hugging Face:

• https://huggingface.co/OliverRen/xAR
• https://huggingface.co/OliverRen/xAR/tree/main

Download examples:

mkdir -p checkpoints

# xAR-B
wget -O checkpoints/xAR-B.pth https://huggingface.co/OliverRen/xAR/resolve/main/xAR-B.pth

# xAR-L
wget -O checkpoints/xAR-L.pth https://huggingface.co/OliverRen/xAR/resolve/main/xAR-L.pth

# xAR-H
wget -O checkpoints/xAR-H.pth https://huggingface.co/OliverRen/xAR/resolve/main/xAR-H.pth

Alternatively, use huggingface-cli:

huggingface-cli download OliverRen/xAR xAR-B.pth --local-dir checkpoints
huggingface-cli download OliverRen/xAR xAR-L.pth --local-dir checkpoints
huggingface-cli download OliverRen/xAR xAR-H.pth --local-dir checkpoints

Model mapping:

• xAR-B.pth -> MODEL_SIZE=base
• xAR-L.pth -> MODEL_SIZE=large
• xAR-H.pth -> MODEL_SIZE=huge

Upstream xAR repository: https://github.com/OliverRensu/xAR

VAE

XYZFlow uses the MAR KL16 VAE tokenizer (kl16.ckpt). A public checkpoint is available on Hugging Face:

• https://huggingface.co/xwen99/mar-vae-kl16
• https://huggingface.co/xwen99/mar-vae-kl16/blob/main/kl16.ckpt

Download example:

wget -O checkpoints/kl16.ckpt https://huggingface.co/xwen99/mar-vae-kl16/resolve/main/kl16.ckpt

Then pass the checkpoint as:

• VAE_PATH=/path/to/kl16.ckpt for training visualizations
• --vae-path /path/to/kl16.ckpt for inference

Data Requirements

XYZFlow student training uses precomputed teacher trajectories, not raw ImageNet images directly. Set TRAJECTORY_DIR to a trajectory directory generated by sample_xyzflow_trajectory_ddp_batch.sh.

ImageNet is still required for:

• training or reproducing the teacher models
• computing FID / IS / Precision / Recall against ImageNet statistics or images
• running the baseline evaluation command in eval.py through --data_path

Download ImageNet-1K / ILSVRC2012 from the official ImageNet website after accepting the dataset terms:

• https://www.image-net.org/challenges/LSVRC/2012/
• https://image-net.org/download-images

Typical files needed for ImageNet-1K experiments are:

• ILSVRC2012_img_train.tar
• ILSVRC2012_img_val.tar
• the ILSVRC2012 development kit / validation labels

Arrange the extracted dataset so that --data_path points to the ImageNet root or training split expected by your evaluation/training command.

ImageNet, generated trajectory data, checkpoints, generated samples, and training outputs are not included. Place them outside the repository or in ignored directories.

Recommended Settings

The main paper explicitly reports the following ImageNet 256x256 setup:

• Teacher checkpoints: xAR-B / xAR-L / xAR-H
• Teacher trajectory generation: 50 denoising steps with CFG 2.3
• Number of precomputed trajectories: 2.5M
• Training hardware: 8 NVIDIA H100 GPUs
• Training iterations: 300K
• Learning rate: 1e-4
• EMA decay: 0.9999
• Student initialization: teacher checkpoint weights
• Next Shortcut inference schedule: 5 -> 4 -> 3 -> 2 patch steps (14 total steps)

The appendix also lists the regression-training configuration as weight decay 0.0, gradient clipping 1.0, and batch size 64. The launch commands below keep the paper-level values where they are directly exposed by this release. Paths and output names are examples.

Generate Teacher Trajectories

CHECKPOINT=checkpoints/xAR-B.pth \
MODEL_SIZE=base \
OUTPUT_ROOT=trajectories_runs \
OUTPUT_PREFIX=xyzflow_base \
NUM_GPUS=8 \
NUM_SAMPLES=2500 \
LABEL_START=0 \
LABEL_END=999 \
BATCH_SIZE=16 \
STEPS_PER_PATCH=50 \
SAVE_STEPS="0 10 20 30 40 50" \
CFG=2.3 \
bash sample_xyzflow_trajectory_ddp_batch.sh

For a quick smoke test:

CHECKPOINT=checkpoints/xAR-B.pth \
MODEL_SIZE=base \
NUM_GPUS=1 \
NUM_SAMPLES=1 \
LABEL_START=0 \
LABEL_END=9 \
BATCH_SIZE=2 \
bash sample_xyzflow_trajectory_ddp_batch.sh

Generated trajectories are organized as:

<trajectory_run>/
  patch_A/step_000/label000/sample0000.pt
  patch_A/step_010/label000/sample0000.pt
  ...
  patch_D/step_050/label999/sampleXXXX.pt

Train XYZFlow

TRAJECTORY_DIR=/path/to/trajectories_runs/xyzflow_base_TIMESTAMP \
CHECKPOINT=checkpoints/xAR-B.pth \
MODEL_SIZE=base \
OUTPUT_DIR=outputs/xyzflow_base \
NUM_GPUS=8 \
BATCH_SIZE=16 \
LR=1e-4 \
CFG=1.0 \
IMAGE_CFG=2.3 \
EMA_DECAY=0.9999 \
MIXED_PRECISION=bf16 \
SAVE_EVERY_STEPS=2000 \
bash train_teacher_student_patchschedule.sh

The paper trains for 300K iterations. This release controls stopping through the launch script/trainer configuration, so set your scheduler or job duration accordingly when reproducing full results.

Enable TensorBoard image previews by setting:

VAE_PATH=/path/to/kl16.ckpt

Export Weights

Training checkpoints contain optimizer state and metadata. Export model weights before inference:

python export_checkpoint.py \
  --input outputs/xyzflow_base/student_teacher_step0300000.pth \
  --output checkpoints/xyzflow_base_ema.pth \
  --weights ema

Use --weights model to export non-EMA weights.

Inference

Example with a 5,4,3,2 patch schedule:

python generate_single_image_patchschedule.py \
  --checkpoint checkpoints/xyzflow_base_ema.pth \
  --vae-path /path/to/kl16.ckpt \
  --arch base \
  --label 0 \
  --cfg 2.3 \
  --per-patch-steps 5,4,3,2 \
  --retain-encoder-layers 6 \
  --retain-decoder-layers 0 \
  --sampler euler \
  --output generated/label000.png \
  --device cuda \
  --no-compile \
  --print-schedule

Evaluation

eval.py evaluates the uniform-step sampling path.

torchrun --nproc_per_node=8 eval.py \
  --model xyzflow_base \
  --resume checkpoints/xAR-B.pth \
  --vae_path /path/to/kl16.ckpt \
  --output_dir outputs/eval_uniform \
  --data_path /path/to/imagenet/train \
  --num_images 50000 \
  --eval_bsz 64 \
  --num_steps 5 \
  --cfg 2.3 \
  --evaluate

Citation

@inproceedings{liu2026xyzflow,
  title={XYZFlow: Scaling Multidimensional Shortcut Flows for Efficient Generative Modeling},
  author={Liu, Jinxiu and Liu, Xuanming and Mei, Kangfu and Wen, Yandong and Liu, Weiyang},
  booktitle={Proceedings of the International Conference on Machine Learning (ICML)},
  year={2026}
}

Acknowledgements

This project builds on the upstream teacher checkpoints. Please also refer to the upstream repository and checkpoints:

• https://github.com/OliverRensu/xAR
• https://huggingface.co/OliverRen/xAR