FAIL: Flow Matching Adversarial Imitation Learning for Image Generation
FAIL employs adversarial imitation learning with pathwise and policy gradient algorithms, improving image generation with only 13,000 samples, surpassing traditional preference methods.
Key Findings
Methodology
FAIL formulates post-training of flow models as an adversarial imitation learning problem, introducing two algorithms: FAIL-PD leverages differentiable ODE solvers for low-variance pathwise gradients, while FAIL-PG uses black-box policy gradients. The framework employs a discriminator to guide the generator without explicit rewards or preference pairs. Experiments fine-tune the FLUX model with only 13,000 expert demonstrations, achieving significant improvements on benchmarks like UniGen-Bench (from 61.61 to 73.70) and DPG-Bench (to 87.32). The approach effectively mitigates reward hacking and enhances stability, generalizing to discrete image and video generation scenarios.
Key Results
- FAIL-PD improves the base model from 61.61 to 73.70 on UniGen-Bench, outperforming FLUX Ultra, and reaches 87.32 on DPG-Bench, with only 13,000 samples. It surpasses RLHF and DPO in stability and quality, avoiding reward model collapse.
- Compared to preference optimization, FAIL leverages continuous discriminator signals, providing richer feedback and better sample efficiency. It demonstrates superior long-term stability, with FAIL-PD maintaining performance beyond 2000 steps, unlike FAIL-PG which collapses after 450 steps.
- The methods show robustness across different discriminator architectures, with the flow matching backbone yielding the best results due to its rich semantic features. Combining FAIL with reward models further prevents reward hacking, achieving the highest overall scores.
Significance
This work advances the application of adversarial imitation learning in high-dimensional image generation, especially under limited data conditions. It addresses the core challenge of aligning generated distributions with high-quality targets without relying on costly preference data or reward models. The proposed algorithms enable stable, efficient training, opening pathways for scalable, data-efficient content creation tools in industry. The ability to generalize to discrete and video domains broadens its impact, promising a new paradigm for robust, reward-free generative modeling.
Technical Contribution
The paper introduces FAIL, a novel framework that combines pathwise derivatives and policy gradients within a flow matching context, leveraging the differentiability of ODE solvers for low-variance gradient estimation. It innovates by formalizing the post-training as an adversarial imitation game, integrating continuous discriminator signals to guide the generator. The algorithms extend the scope of flow models, enabling stable, high-fidelity generation with minimal data, and provide theoretical insights connecting DPG and GANs. The approach also offers practical solutions for black-box and discrete scenarios, broadening the applicability of adversarial imitation learning.
Novelty
This is the first work to formalize flow matching model post-training as an adversarial imitation learning problem, introducing two algorithms—FAIL-PD and FAIL-PG—that address different computational settings. It departs from reliance on explicit reward models, instead utilizing a discriminator’s continuous signals for training stability and efficiency. The integration of differentiable ODE solvers for low-variance pathwise gradients represents a significant innovation, enabling high-quality image generation with limited data.
Limitations
- FAIL-PD relies heavily on the differentiability of ODE solvers, which can be computationally intensive and sensitive to numerical stability, limiting scalability.
- FAIL-PG, while flexible, exhibits instability over long training horizons, requiring additional regularization to prevent collapse.
- The current evaluation is primarily on image datasets; performance on more complex, multi-modal, or video generation tasks remains to be validated.
Future Work
Future directions include optimizing the efficiency of pathwise gradient computation, extending the framework to multi-modal and video domains, and developing hybrid algorithms that combine the stability of FAIL-PD with the speed of FAIL-PG. Further research will explore integrating more sophisticated discriminator architectures and reward models to enhance robustness, as well as scaling the approach to larger datasets and real-world applications.
AI Executive Summary
Recent advances in generative modeling, such as flow matching and diffusion models, have revolutionized image synthesis, enabling high-fidelity, diverse outputs. However, the post-training phase—aligning models with high-quality targets—remains challenging, especially under limited data conditions. Traditional supervised fine-tuning (SFT) effectively mimics expert demonstrations but suffers from distribution shift, leading to degraded performance on unseen states. Preference optimization methods like RLHF and DPO address this by leveraging scalar rewards or preference pairs, yet they are costly and prone to reward hacking.
Ma et al. propose FAIL, a novel adversarial imitation learning framework that reformulates flow model post-training as a minimax game between a generator policy and a discriminator. Unlike prior methods, FAIL does not require explicit reward functions or preference pairs. Instead, it employs two algorithms: FAIL-PD, which exploits the differentiability of ODE solvers to compute low-variance pathwise gradients, and FAIL-PG, a black-box policy gradient method suitable for discrete or computationally constrained scenarios. The framework is validated by fine-tuning the FLUX model with only 13,000 expert demonstrations, achieving significant improvements on benchmark datasets. For instance, FAIL-PD boosts performance on UniGen-Bench from 61.61 to 73.70, surpassing existing models.
Experimental results demonstrate that FAIL not only outperforms traditional preference optimization in data efficiency and stability but also generalizes effectively to discrete image and video generation tasks. The approach mitigates reward hacking and enhances training robustness, offering a scalable solution for high-quality content synthesis with limited supervision. Its ability to integrate seamlessly with reward models further broadens its applicability, providing a flexible and stable alternative to existing methods.
This work marks a significant step toward stable, data-efficient, reward-free generative modeling, with broad implications for AI content creation, virtual reality, and multimedia applications. Future research will focus on improving computational efficiency, extending to multi-modal data, and exploring hybrid algorithms to combine the strengths of pathwise and policy gradient approaches, pushing the boundaries of scalable, robust AI generation systems.
Deep Analysis
Background
The evolution of image generation has transitioned from GANs to diffusion and flow-based models, emphasizing high-quality, diverse synthesis. Flow matching models, such as those introduced by Lipman et al. and Liu et al., leverage continuous transformations for stable training. Despite预训练阶段的突破，后训练仍依赖偏好数据和奖励模型，存在偏差崩溃和数据稀缺的难题。偏好优化（如RLHF）虽有效，但成本高昂，限制了其应用范围。近年来，基于对抗学习的方法逐渐兴起，试图在无需偏好对的情况下实现高效对齐，但仍面临训练不稳定和样本效率低的问题。
Core Problem
核心难题在于如何在少样本、无偏好数据的环境下，实现模型输出分布与高质量目标的有效对齐。传统微调易受分布偏移影响，偏好优化依赖昂贵偏好对，训练过程不稳定。现有对抗方法虽能缓解部分问题，但缺乏结构化的梯度信息，导致训练不稳定和崩溃风险高。如何利用模型的连续性和判别器的丰富反馈，设计稳定、高效的后训练算法，是亟待解决的关键。
Innovation
主要创新包括：1）将流匹配模型的后训练问题转化为对抗模仿学习，避免显式奖励依赖；2）提出FAIL-PD算法，利用ODE的可微性实现低方差路径梯度，提升训练稳定性；3）开发FAIL-PG算法，适应离散和计算受限场景，提供黑盒策略。结合判别器的连续信号，有效缓解偏差崩溃，显著提升少样本训练效果。引入正则化策略，增强训练稳定性，推动模型在有限示范下的高质量生成。
Methodology
- �� 构建生成器策略πθ与判别器Dω的对抗框架，通过极大化判别器区分生成样本与专家样本的能力，目标是最小化两者的分布差异。• FAIL-PD利用ODE的可微性，通过单步去噪近似，反向传播判别器梯度，更新策略参数。• FAIL-PG采用奖励信号，利用策略梯度优化，适应离散和复杂场景。• 结合路径梯度和黑盒奖励，设计多样化训练流程，确保模型在少样本条件下的高效学习。• 采用预训练视觉模型作为判别器基础，提升特征表达能力，增强判别效果。
Experiments
使用Gemini 3 Pro作为专家示范，构建13,000个prompt-图像对，模拟少样本环境。模型基于FLUX架构，判别器采用预训练视觉模型。训练过程中，比较FAIL与RLHF、DPO等偏好优化方法，评估指标包括prompt跟随和审美质量。采用多GPU分布式训练，调优超参数，验证算法稳定性和泛化能力。还进行长时间训练的稳定性分析和不同判别器架构的影响研究。
Results
FAIL在UniGen-Bench将性能从61.61提升至73.70，超越FLUX Ultra，DPG Bench达87.32。少样本条件下，模型在prompt跟随和审美指标上表现优异，验证了其数据效率。与偏好优化方法相比，FAIL避免了偏差崩溃，利用连续判别信号实现更丰富反馈。长时间训练中，FAIL-PD表现出更高的稳定性，避免模型崩溃，验证路径梯度的优势。
Applications
该方法适用于需要高质量图像生成的场景，如内容创作、广告设计和虚拟现实。尤其在偏好数据难以获取或有限示范条件下，能快速实现高质量输出。未来还可扩展到视频和多模态生成，推动工业界智能内容生成的普及。
Limitations & Outlook
当前方法对ODE求解器依赖较强，计算成本高，且在大规模应用中存在性能瓶颈。FAIL-PG在长时间训练中表现出不稳定，需额外正则化。模型在多模态和视频生成中的适应性仍需验证，未来需优化算法效率和稳定性。
Plain Language Accessible to non-experts
想象你在一个工厂里，工厂的目标是生产高质量的产品。传统方法就像让工人一遍遍模仿专家的操作，但如果遇到新情况，工人可能会偏离正确路线。现在，有一种智能助手，它可以通过观察专家的操作，学习如何调整工厂的机器，使生产线变得更智能、更高效。这个助手不用事先知道所有的奖励，只通过和判别器的“对话”不断改进自己。它像一个聪明的学生，通过不断练习和老师的反馈，逐渐掌握了最优的生产方法。这个过程既快又稳，不需要大量的示范，也能适应不同的生产任务。最终，工厂的产品质量大幅提升，效率也更高了。这就像论文中的FAIL方法，用对抗训练让模型在少量示范下学会生成高质量图片，避免了繁琐的奖励设计，变得更智能、更稳定。
ELI14 Explained like you're 14
想象你在学校学画画，老师给你一些漂亮的画作作为示范。你试着模仿，但有时候你会偏离老师的风格，画得不够好。这时候，如果有个聪明的朋友，他会偷偷告诉你哪里画得像老师，哪里还可以改进。这个朋友不用告诉你具体的奖励，只是通过观察你的画和老师的画，告诉你哪些地方做得更好。你不断调整自己的画风，变得越来越像老师的风格，而且还学得很快。这个朋友就像论文中的判别器，它帮助你在没有明确奖励的情况下，快速学会高质量的画画技巧。通过这种对抗的学习方式，你可以用很少的示范，画出令人满意的作品，而且过程很稳定，不会突然崩溃。这就是FAIL的核心思想，用对抗训练让模型在少量示范下变得更聪明、更稳定，能生成漂亮的图片。
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.