MGFlow

Model checkpoints, reference distributions, and evaluation assets for MGFlow, presented in Unifying Distributional Training for One-Step Visual Generation.

The release covers ImageNet class-conditional generation at 256×256 with JiT and pMF (B/L/H), and text-to-image generation at 512×512 with FLUX.2 [klein] 4B. Both generate images in one step after MGFlow post-training.

The T2I checkpoints use the MGFlow-T2I reference dataset, which contains synthetic COCO and GenEval images paired with their prompts. The dataset link in the metadata applies to this T2I branch.

Project page: https://shihaoyang0423.github.io/MGFlow-website

Directory	Contents
Checkpoints/ImageNet/Base	6 pretrained initialization checkpoints: JiT/pMF-B, L, H.
Checkpoints/ImageNet/Post-trained	12 checkpoints: JiT/pMF-B, L, H, each with W2 (K=1+4) and KL (K=1+4+16).
Checkpoints/T2I	2 FLUX.2 [klein] 4B checkpoints: image-only and joint image–text.
Evaluation/ImageNet	7 evaluation statistics files.
Reference	Gaussian reference distributions and frozen encoder weights.

License

See LICENSE.md for asset-specific licenses and upstream notices.

Citation

If you use MGFlow or these assets, please cite:

@misc{zhang2026unifyingdistributionaltrainingonestep,
  title={Unifying Distributional Training for One-Step Visual Generation},
  author={Chi Zhang and Shi Haoyang and Yueyi Liu and Ruichuan An and Junkang Zhou and Chang Li and Xiuyuan Lu and Yichi Zhang and Bo Wang and Yuhang Wu and Sen Cui and Miao Liu},
  year={2026},
  eprint={2609.35763},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2609.35763}
}