FAIL: Flow Matching Adversarial Imitation Learning for Image Generation

TL;DR

FAIL introduces a novel adversarial imitation learning framework for image generation that aligns policies with expert demonstrations without explicit rewards, improving performance and robustness in various generative tasks.

Contribution

The paper presents FAIL, a new adversarial imitation learning method for image generation that avoids reward modeling and preference pairs, with algorithms suitable for different computational settings.

Findings

FAIL achieves competitive prompt following performance.

The framework generalizes to image and video generation.

Fails acts as a regularizer to reduce reward hacking.

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. Code and data are available at https://github.com/HansPolo113/FAIL.