Search papers, labs, and topics across Lattice.
The paper introduces Flow Matching Adversarial Imitation Learning (FAIL), a novel approach to fine-tuning flow matching models for image generation by framing the alignment with a target distribution as an imitation learning problem. FAIL leverages adversarial training to minimize the divergence between the policy and expert demonstrations, avoiding the need for explicit rewards or pairwise comparisons. The authors demonstrate that FAIL achieves competitive performance on prompt following and aesthetic benchmarks with limited demonstrations, and also show its effectiveness in discrete image/video generation and as a regularizer against reward hacking.
Ditch the costly preference pairs and reward modeling: FAIL leverages adversarial imitation learning to fine-tune flow matching models, achieving state-of-the-art image generation with minimal data.
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. Code and data are available at https://github.com/HansPolo113/FAIL.