FAIL: Flow Matching Adversarial Imitation Learning for Image Generation
Yeyao Ma, Chen Li, Xiaosong Zhang, Han Hu, Weidi Xie
Abstract
Post-training of flow matching models-aligning the output distribution with a high-quality target-is mathematically equivalent to imitation learning. While Supervised Fine-Tuning mimics expert demonstrations effectively, it cannot correct policy drift in unseen states. Preference optimization methods address this but require costly preference pairs or reward modeling. We propose Flow Matching Adversarial Imitation Learning (FAIL), which minimizes policy-expert divergence through adversarial training without explicit rewards or pairwise comparisons. We derive two algorithms: FAIL-PD exploits differentiable ODE solvers for low-variance pathwise gradients, while FAIL-PG provides a blackbox alternative for discrete or computationally constrained settings. Fine-tuning FLUX with only 13,000 demonstrations from Nano Banana pro, FAIL achieves competitive performance on prompt following and aesthetic benchmarks. Furthermore, the framework generalizes effectively to discrete image and video generation, and functions as a robust regularizer to mitigate reward hacking in reward-based optimization. Code and data are available at https://github.com/ HansPolo113/FAIL
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Builds on29
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Training language models to follow instructions with human feedbackLong Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida et al.NeurIPS 2022 · 24,707 citations
- Denoising Diffusion Implicit ModelsJiaming Song, Chenlin Meng, Stefano ErmonICLR 2021 · 11,743 citations
- Direct Preference Optimization: Your Language Model is Secretly a Reward ModelRafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D. Manning et al.NeurIPS 2023 · 10,924 citations
- Scaling Rectified Flow Transformers for High-Resolution Image SynthesisPatrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari et al.ICML 2024 · 3,620 citations
Related papers
- PISA Experiments: Exploring Physics Post-Training for Video Diffusion Models by Watching Stuff DropChenyu Li, Oscar Michel, Xichen Pan, Sainan Liu et al.ICML 2025
- Unpaired Visual Editing with Self-Consistent Flow MatchingYoad Tewel, Yuval Atzmon, Gal Chechik, Lior WolfICML 2026
- PC-Flow: Preference Alignment in Flow Matching via ClassifierShaomeng Wang, He Wang, Longquan Dai, Jinhui TangAAAI 2026
- Online Reward-Weighted Fine-Tuning of Flow Matching with Wasserstein RegularizationJiajun Fan, Shuaike Shen, Chaoran Cheng, Yuxin Chen et al.ICLR 2025
- Test-Time Reinforcement Learning for Flow MatchingJili Chen, Changqin Huang, Qionghao Huang, Yaxin Tu et al.ICML 2026
