ReflexFlow: Rethinking Learning Objective for Exposure Bias Alleviation in Flow Matching

TL;DR

ReflexFlow introduces a reflexive refinement of the Flow Matching learning objective to dynamically correct exposure bias, improving generative quality across multiple datasets by addressing model generalization and low-frequency content issues.

Contribution

The paper proposes ReflexFlow, a novel method with Anti-Drift Rectification and Frequency Compensation components, to effectively mitigate exposure bias in Flow Matching models.

Findings

Achieves a 35.65% reduction in FID on CelebA-64.

Outperforms prior methods in exposure bias mitigation.

Compatible with all Flow Matching frameworks.

Abstract

Despite tremendous recent progress, Flow Matching methods still suffer from exposure bias due to discrepancies in training and inference. This paper investigates the root causes of exposure bias in Flow Matching, including: (1) the model lacks generalization to biased inputs during training, and (2) insufficient low-frequency content captured during early denoising, leading to accumulated bias. Based on these insights, we propose ReflexFlow, a simple and effective reflexive refinement of the Flow Matching learning objective that dynamically corrects exposure bias. ReflexFlow consists of two components: (1) Anti-Drift Rectification (ADR), which reflexively adjusts prediction targets for biased inputs utilizing a redesigned loss under training-time scheduled sampling; and (2) Frequency Compensation (FC), which reflects on missing low-frequency components and compensates them by reweighting the loss using exposure bias. ReflexFlow is model-agnostic, compatible with all Flow Matching frameworks, and improves generation quality across datasets. Experiments on CIFAR-10, CelebA-64, and ImageNet-256 show that ReflexFlow outperforms prior approaches in mitigating exposure bias, achieving a 35.65% reduction in FID on CelebA-64.