FAIL: Flow Matching Adversarial Imitation Learning for Image Generation

Framing flow matching post-training as imitation learning: without preference pairs or explicit reward models, bring the generation distribution closer to expert demonstrations.

• Scenario: the authors argue that pretraining covers a broad distribution, while post-training aligns to a high-quality preferred distribution to improve aesthetics and prompt following. The experimental setting is one expert image per prompt.
• Existing shortcomings: SFT corresponds to behavioral cloning. The authors state it is data-efficient but drifts and accumulates errors in states outside the expert demonstrations. RLHF and DPO handle drift but rely on preference data and reward models.
• Observation: adversarial imitation learning uses a discriminator to distinguish expert samples from generated samples, minimizing the Jensen-Shannon divergence via minimax, providing dense signals without explicitly modeling rewards or pairwise comparisons.
• Proposal: FAIL (Flow Matching Adversarial Imitation Learning). White-box FAIL-PD backpropagates through the differentiable ODE; black-box FAIL-PG is for discrete generation or settings where backpropagation is too costly. The authors also use it as a regularizer for reward optimization.

The introduction and Related Work divide related work into preference optimization, adversarial training on diffusion models, and imitation learning.

Preference optimization

• SFT / behavioral cloning: Pomerleau (1991a) directly imitates experts with regression. The authors state that SFT on curated data (e.g., Emu, Dai et al., 2023) improves generation quality but often misses fine-grained human preferences and suffers from distribution shift.
• RLHF with explicit rewards: reward gradient methods (Xu et al., 2023; Clark et al., 2023; Shen et al., 2025) guide the model with reward score gradients; policy gradient methods (Black et al., 2023; Xue et al., 2025b; Liu et al., 2025) treat denoising as sequential decision-making to maximize expected reward. The authors state that explicit reward modeling comes with instability and reward hacking.
• DPO: Rafailov et al. (2023) and Diffusion-DPO (Wallace et al., 2024) optimize directly on preference pairs, bypassing explicit reward models.

Adversarial training on diffusion models

• Objective augmentation: DiffusionGAN (Wang et al., 2022), DDGAN (Xiao et al., 2021), UFOGen (Xu et al., 2024), and SiDDM (Gong et al., 2023) use discriminators on generated samples or intermediate diffusion states, supplementing the diffusion objective with a GAN loss.
• Adversarial distillation: ADD (Sauer et al., 2024b), LADD (Sauer et al., 2024a), and DMD2 (Yin et al., 2024) use adversarial objectives to distill a multi-step diffusion teacher into a student with fewer steps.

Imitation learning

• BC and IRL: BC (Pomerleau, 1991b; Ross et al., 2011) directly regresses to imitate experts; IRL (Ziebart et al., 2008; Abbeel & Ng, 2004) jointly infers policy and reward.
• AIL / GAIL: Ho & Ermon (2016) and Fu et al. (2017b) use adversarial training to match the state-action distribution of the agent and the expert. The authors state that Orsini et al. (2021) and Kostrikov et al. (2018) improved GAIL, and GAIL has been applied in multiple domains (Wulfmeier et al., 2024; Ye et al., 2025; Mao et al., 2025).

Baselines and benchmarks

• Baseline methods: comparisons against expert data include SFT, RLHF-PD, RLHF-PG, and Online DPO; when used with reward models, comparisons include Reward Gradient and FPO. Policy updates also borrow from FPO (McAllister et al., 2025), PPO (Schulman et al., 2017), and GRPO (Shao et al., 2024).
• Benchmarks: UniGen-Bench, DPG-Bench, Alchemist (reporting HPSv3 and UnifiedReward), HPDv3, and VBench.

The authors frame post-training as adversarial imitation, using FAIL to connect AIL to the continuous-time dynamics of flow matching, thereby not relying on preference pairs or reward models. FAIL-PG replaces the TRPO step with FPO in the GAIL framework.

FAIL formulates post-training as a minimax between policy π_θ and discriminator D_ω: generated samples must be indistinguishable from the expert distribution p_E. Updates are divided into the white-box path derivative FAIL-PD and the black-box policy gradient FAIL-PG.

Problem Setting

• Objective: The paper writes

min_θmax_ω𝔼_{x∼ pE}[log D_ω(x)]+𝔼_x∼πθ[log(1−D_ω(x))] i.e., the discriminator distinguishes expert from policy samples, and the policy reduces the divergence between them. In Algorithms 1–2, the discriminator output first passes through σ and then takes log, corresponding to the logit form.

• Dynamics: Samples are obtained by integrating the vector field v_θ from noise ϵ∼N(0,I). The authors state that the "environment" here is a deterministic ODE solver, so gradients can be backpropagated without learning a separate critic as in an unknown environment.
• Data setting: Each prompt is paired with only one expert image, to address cases where preference pairs are difficult to collect.

FAIL-PD

• Path derivative: The discriminator evaluates the generated x_0 and serves as a differentiable reward. Euler or RK4 is differentiable, and gradients propagate to θ via the chain rule. The authors relate this to deterministic policy gradients (Silver et al., 2014) and call it a low-variance, unbiased DDPG limit for imitation learning, inheriting GAN's local convergence properties.
• Single-step approximation: It does not unroll the entire trajectory for each gradient. Sample t∼U(0,1), add noise as x_t=(1-t)x_0+tϵ, then do one denoising step to estimate the clean sample x′_0; Δt is the inference step size. The authors state this only assumes local linearity on that step. Appendix B derives this formula from the interpolation path and tests it with reconstruction error.
• Policy update: Minimize log σ(D(x′_0)) with respect to x′_0, i.e., the gradient is −∇_θ log σ(D_ω(x′_0)).

FAIL-PG

• Scalar reward: It does not backpropagate through the dynamics. The reward is r(x)=−log(1−σ(D_ω(x))), reinforcing high-reward samples. The authors state this corresponds to GAIL but applied to continuous flow matching.
• Likelihood ratio: It does not convert the ODE into an SDE to compute an exact likelihood. Using FPO, the likelihood ratio between the old and current policies is approximated by the CFM loss difference

r(θ)=exp(L_CFM(θ_old,x)−L_CFM(θ,x)) and then a PPO-style clipped objective is maximized. The KL to the reference policy is likewise written as the difference of two CFM losses. Appendix A uses an ELBO approximation to show that constant terms cancel.

• Within-group normalization: Advantages are standardized within a group of G samples using mean-standard deviation normalization following GRPO. The algorithm first updates the discriminator, then performs several epochs of KL-constrained FPO with the old policy. The paper does not give the clipping coefficient ϵ or the number of inner epochs E; the KL coefficient β in Appendix C is 0.05.

Stabilization and Discriminator

• Mixed imitation: Following Ren et al. (2024), the batch mixes online policy samples with expert demonstrations; expert samples are treated as perfect rollouts to reduce exploration in large state spaces.
• Initialization: Behavioral cloning is used for cold start. The policy is frozen during the discriminator warm-up phase; the default warm-up in ablations is 25 steps.
• Three backbones: VFM adds a projection head on intermediate features from CLIP, DINO, etc.; the authors state it has no native text input and smaller capacity. The FM backbone loads a pretrained flow model with all parameters trainable and attaches a discriminator head on intermediate features. VLM looks at both the image and the prompt, taking the hidden state of the last token; the vision encoder is frozen, and the LLM and discriminator head are fine-tuned. The main experiments use Qwen3-VL-2B-Instruct. Appendix C: the VFM/FM head is LayerNorm plus a linear layer with spectral normalization; the VLM takes the EOS token hidden state, then passes it through one linear mapping.
• The authors' division of labor for the two updates: FAIL-PD provides a directional gradient like ∇_x D(x), which the authors believe better preserves manifold structure and is compatible with CFG; FAIL-PG is a mode-seeking scalar reward that is faster early on, has higher variance, and may exhibit model collapse over the long range. The former is used for differentiable white-box settings, the latter for black-box or discrete tokens.

Experimental Setup

• Policy: The main image experiment uses FLUX.1-dev; discrete images use Xomni; video uses student Wan2.1-1.3B.
• Expert data: Prompts come from a subset of HPDv3 and train-en-short from UniGen-Bench. The abstract calls the expert Nano Banana pro, Section 3.1 writes Gemini Pro 3, and the introduction writes Gemini 3 Pro; one image per prompt, 13,000 pairs after filtering. Video draws 8,192 prompts from VidProM, with expert Wan2.2-A14B, resolution 480×480, 81 frames.
• Training: AdamW, global batch 128, 32 NVIDIA H20. The effective batch per prompt is 3 policy rollouts plus 1 expert image. Unless otherwise stated, ablations are 512×512, 1 epoch (400 iterations), including 25 steps of discriminator warmup. Learning rates are listed by task in Table 11. FAIL-PG's β is 0.05.
• Discriminator: The main experiment uses Qwen3-VL-2B-Instruct; also Dinov3b and FLUX.1-dev backbones.
• Metrics and evaluation: UniGen's en-short has 600 items; ablations use Qwen-72B, system-level uses Gemini 2.5 Pro API. DPG-Bench has 1,065 items. Preferences use Alchemist's 3,350 items, reporting HPSv3 and UnifiedReward-2.0-qwen3vl-8b. UniGen and DPG use 4 images per prompt, Alchemist uses 1. Ablation inference uses 28 steps, 512×512, CFG off. System-level FAIL samples are still 512×512 but use optimal CFG; the paper does not give that CFG scale.
• Repetition: The paper does not report random seeds or number of repetitions. The protocols in Table 1 and Tables 6–8 differ, so scores cannot be directly compared.

Main Results

Comparisons with the same expert data and CFG off are in Table 1. There, UniGen belongs to the ablation protocol (Qwen-72B).

Swipe horizontally to read the table

The table omits RLHF-PD at 50 steps: UniGen 59.76, DPG 81.65, HPSv3 10.18, URScore 3.3406 (Table 1).

• Reward hacking: The authors say RLHF training reward rises while ground metrics worsen; only 50-step RLHF-PD exceeds SFT. 400-step RLHF-PD and 50-step RLHF-PG are not higher than SFT on several metrics.
• Relative to Online DPO: The authors view Online DPO as a special case of FAIL-PG where the discriminator always prefers the expert and the reward is binarized. 400-step FAIL-PG is higher than DPO-Online on all four metrics; FAIL-PD at the same step count is higher only on HPSv3. 800-step FAIL-PD exceeds that DPO on all four metrics. The authors attribute this to dynamic discrimination reducing false negatives and continuous logits being denser than binary labels.
• Step count: The authors say FAIL-PD improves more slowly than FAIL-PG, reaching similar performance only at 800 steps.

Convergence and Stability

• What was measured: Extended training, comparing the two updates with UnifiedReward (Section 3.3, Figure 2). The caption says PG converges quickly but can collapse, while PD has long-term stability; the curve shapes are not visible in the text conversion.
• Results: The authors say FAIL-PG reaches UnifiedReward above 3.43 within 50 steps and collapses after 450 steps; FAIL-PD reaches 3.4759 by 2,000 steps with no collapse seen.
• Interpretation: The authors believe a scalar reward lacks structured gradients, and the long horizon is insufficient to keep the policy on the flow manifold. Appendix E uses KL masking (c=2 unless otherwise stated) to train FAIL-PG to 800 steps: UniGen 65.49, DPG 84.81, HPSv3 11.45, URScore 3.4507 (Table 12).

System-Level Comparison

• Protocol: Same FAIL checkpoint, 512×512, optimal CFG (Section 4). UniGen is instead evaluated by Gemini 2.5 Pro. The paper does not state whether other models in Tables 6–8 were regenerated by the authors under the same protocol, nor the training step count of that checkpoint.
• DPG: In Table 6, FAIL-PD Overall is 87.32, FAIL-PG is 86.42, FLUX.1 Dev is 83.84. The authors say this is 3.48 points above the base model and comparable to Seedream; in the table, Seedream 3.0 is 88.27 and Qwen-Image is 88.32.
• UniGen and HPDv3: In Table 7, FAIL-PD Overall is 73.70, FAIL-PG is 72.64, FLUX.1 Dev is 61.30. The text writes the baseline as 61.61, which matches Janus-Pro-7B's Overall in the table and does not match FLUX.1 Dev; the table is taken as authoritative here. The authors say it exceeds FLUX.1 Ultra (70.67) and Imagen-3.0 (71.85); in the table, Qwen-Image is 78.81 and Seedream 3.0 is 78.98. Text goes from baseline 32.18 to FAIL-PD 52.87, with Qwen-Image at 76.14; the authors say post-training cannot make up spelling ability not established in pretraining. In Table 8, FAIL-PD's All is 11.28 and FLUX.1 Dev is 10.43; in the Science category FAIL-PD is 9.59, below FLUX.1 Dev's 11.24. The paper does not explain this category decline.

Reward Combination and Cross-Modality

• Regularization: In Table 3, FAIL-PD alone has the same four metrics as the 400-step row in Table 1. Reward Gradient alone has HPSv3 12.21, UniGen 58.50, DPG 82.15. After addition, HPSv3 is 12.94 and URScore 3.4832, the highest of these two metrics in that table; UniGen 60.88 is still below FAIL-PD alone. FPO+FAIL-PG has DPG 84.70, and UniGen 64.27 is below FAIL-PG alone at 64.75. The authors say dynamic distribution matching can mitigate reward hacking on static rewards. Appendix G.2 says optimizing HPSv3 alone shows oversaturation and high-frequency artifacts; Figure 5 is not visible.
• Discrete images: Table 4 reports only FAIL-PG. Xomni's UniGen goes from 28.32, SFT 41.36, to 54.24; DPG 80.94, HPSv3 9.91, URScore 3.3591. The discriminator is DINOv3b. The authors say online samples mitigate the next-token train-inference gap.
• Video: Table 5. Wan2.1 base Overall / Quality / Semantic is 68.53 / 76.43 / 36.93, SFT is 73.02 / 79.71 / 46.27, FAIL-PG is 75.71 / 81.70 / 51.75, and FAIL-PD's Semantic is 49.07. The authors attribute the relatively small Quality gain to no CFG during rollout.

Other Ablations and Analyses

• Discriminator (Table 2, FAIL-PD): HPSv3 is Dinov3 10.28, Qwen3vl 10.82, FLUX.1 12.35; UniGen is 60.65, 62.17, 62.70.
• Data scale (Table 13, 800-step FAIL-PD UniGen): 4K 60.04, 8K 62.29, 13K 63.38.
• Cold start (Table 15, 800 steps): without cold start UniGen 55.23, below SFT's 58.22; with cold start it is 63.38. The authors say running FAIL-PG or Online DPO directly without cold start can collapse within about 100 steps.
• Distribution distance (Table 16, MMD×10^3 / SWD×10^3): FLUX 18.26 / 7.10, SFT 9.60 / 5.19, FAIL-PD 2000 steps 2.24 / 2.76. Features are SigLIPv2 so400m, 1,024 expert samples.
• Single-step approximation (Table 9): interval [0, 0.5] has PSNR 35.17, LPIPS 0.030; [0.5, 1] has 31.13, 0.115; full interval 32.55, 0.083.
• Per-step time (Table 14, seconds): FAIL-PD total 52.73, of which policy update 12.06; FAIL-PG total 91.85, of which policy update 50.08.

Limitations and future directions noted by the authors

• Optimization sensitivity: despite stabilization strategies, the adversarial objective remains sensitive to hyperparameters and model design (Section 6).
• Data scale: experiments are validated only on 13K; the authors leave FAIL's scaling as an open question (Section 6).
• Base model upper bound: text rendering still has gaps. The authors state that alignment methods can only refine capabilities already established by pretraining, and cannot synthesize basic capabilities absent from the pretraining stage (Section 6).

Experimental coverage

• The policies with reported fine-tuning results are FLUX.1-dev, Xomni, and Wan2.1-1.3B.
• Image expert data is a filtered set of 13,000 single-image pairs; Appendix F additionally provides 4K and 8K. Video uses 8,192 VidProM prompts, with Wan2.2-A14B as the expert model.
• Discriminator experiments include Dinov3b, the FLUX.1-dev backbone, and Qwen3-VL-2B-Instruct. The reward-combined setting uses frozen HPSv3, as well as Reward Gradient and FPO.
• Evaluation includes UniGen-Bench (600 prompts; Qwen-72B for ablations, Gemini 2.5 Pro for system-level), DPG-Bench (1,065 prompts), Alchemist (3,350 prompts), HPDv3, and VBench. Ablations turn off CFG; system-level FAIL samples use optimal CFG, and the resolution is 512×512 throughout. Section 4 does not state the number of training steps for this checkpoint.
• The paper does not report the number of repetitions or random seeds, nor does it state whether the comparison models in Tables 6–8 were rerun under the same generation protocol.

• Positioning: FAIL turns flow matching post-training into adversarial imitation, using a discriminator to match the expert distribution, thereby not collecting preference pairs and not training a separate reward model.
• Method: FAIL-PD backpropagates the discriminator signal through the differentiable ODE, using a single-step denoising approximation during training; FAIL-PG turns the discriminator logit into a reward, then updates with FPO, GRPO, and reference policy KL. Training first does behavior cloning, and warms up the discriminator, mixing in expert samples.
• Same protocol: On FLUX.1-dev with 13,000 single-image expert data, in Table 1 (CFG off, UniGen scored by Qwen-72B), 400-step FAIL-PG's UniGen is 64.75, Online DPO is 62.83. RLHF-PD at 400 steps has UniGen 57.97, lower than SFT's 58.22.
• System level: With optimal CFG, 512×512, FAIL-PD's UniGen Overall is 73.70 (Table 7), DPG Overall is 87.32 (Table 6), HPDv3 All is 11.28 (Table 8). Table 7's Text is 52.87, Qwen-Image is 76.14. The baseline Overall 61.61 written in the text is inconsistent with FLUX.1 Dev's 61.30 in the table.
• Authors' conclusion: FAIL-PG converges faster but collapses after about 450 steps; at 2,000 steps FAIL-PD's UnifiedReward is 3.4759. After adding the HPSv3 reward gradient, Table 3's HPSv3 is 12.94. There are also improvements on Xomni and Wan2.1-1.3B. The authors say adversarial imitation can connect SFT and reinforcement learning, and constrain reward speculation when used with an explicit reward, but it is still constrained by the base pretraining.