FAIL: Flow Matching Adversarial Imitation Learning for Image Generation
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 black-box 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.
Lay Summary
We formulate the post-training of flow matching image generation models as an imitation learning problem, where the expert is a high quality dataset. We proposes an approach called flow matching adversarial imitation learning (FAIL) and proposes some concrete algorithms to fine tune flow matching models using this approach.