Flow-Factory: A Unified Framework for Reinforcement Learning in Flow-Matching Models

Reinforcement learning has emerged as a promising paradigm for aligning diffusion and flow-matching models with human preferences, yet practitioners face fragmented codebases, model-specific implementations, and engineering complexity. We introduce Flow-Factory, a unified framework that decouples algorithms, models, and rewards through through a modular, registry-based architecture. This design enables seamless integration of new algorithms and architectures, as demonstrated by our support for GRPO, DiffusionNFT, and AWM across Flux, Qwen-Image, and WAN video models. By minimizing implementation overhead, Flow-Factory empowers researchers to rapidly prototype and scale future innovations with ease. Flow-Factory provides production-ready memory optimization, flexible multi-reward training, and seamless distributed training support. The codebase is available at https://github.com/X-GenGroup/Flow-Factory.

💡 Research Summary

The paper presents Flow‑Factory, a unified and extensible software framework designed to simplify and accelerate reinforcement‑learning (RL) fine‑tuning of flow‑matching and diffusion models. Recent work has shown that RL can align these powerful generative models with human preferences, but practical adoption is hampered by three intertwined challenges: (1) Fragmented codebases – each RL algorithm (e.g., Flow‑GRPO, Dance‑GRPO, DiffusionNFT, AWM) typically ships with its own SDE/ODE implementation and tightly coupled model‑specific logic, making it difficult to reuse code across different backbone architectures; (2) Training‑time memory bottlenecks – large multimodal models contain frozen components (text encoders, VAE decoders, CLIP backbones) that occupy substantial GPU memory even when only the transformer core is trainable; (3) Limited reward flexibility – existing frameworks support only pointwise reward models, whereas newer methods require groupwise ranking rewards or the combination of several heterogeneous reward signals.

Core Design
Flow‑Factory tackles these problems through three orthogonal design pillars:

1. Registry‑Based Component Decoupling – The framework defines four abstract base classes: BaseAdapter (model I/O, encoding, checkpointing), BaseTrainer (sampling, advantage computation, optimizer steps), BaseRewardModel (pointwise or groupwise scoring), and SDESchedulerMixin (stochastic sampling with log‑probability). Concrete implementations register themselves in a global registry, and a YAML configuration file selects and composes any combination of model, trainer, reward, and scheduler at runtime. This reduces the integration complexity from O(M × N) (M models × N algorithms) to O(M + N).

2. Preprocessing‑Based Memory Optimization – Before training, all condition embeddings (prompt embeddings, pooled embeddings, VAE latents) are pre‑computed and cached on disk. During training, only the trainable transformer resides on the GPU; frozen encoders are completely off‑loaded. The approach cuts peak per‑GPU memory by ~13 % (≈8 GB on an H200) and eliminates redundant forward passes, yielding a 1.74× speed‑up per training step (144 s → 82 s in the authors’ experiments).

3. Flexible Multi‑Reward System – Flow‑Factory unifies pointwise and groupwise reward interfaces, introduces a MultiRewardLoader that deduplicates identical reward models across configurations, and supports configurable advantage aggregation (simple weighted sum or GDPO‑style per‑reward normalization). This enables algorithms such as Pref‑GRPO (ranking‑based reward) and multi‑objective setups that blend Text‑Rendering, PickScore, or other signals without extra engineering.

Supported Algorithms
The framework implements several state‑of‑the‑art RL‑for‑flow methods, all sharing the same adapter and reward interfaces:

• Flow‑GRPO – Extends deterministic ODE sampling with stochastic noise injection, providing a tractable log‑probability for policy‑gradient updates. Different SDE dynamics (Flow‑SDE, Dance‑SDE, CPS, pure ODE) are selectable via a single scheduler parameter.
• MixGRPO – Applies SDE updates only on a few timesteps while using ODE elsewhere, reducing compute cost while preserving performance.
• GRPO‑Guard – Introduces timestep‑dependent loss re‑weighting to mitigate the “negative‑bias” problem of pure SDE formulations.
• DiffusionNFT – Optimizes a contrastive objective directly on the forward flow, eliminating the need for likelihood estimation; it is solver‑agnostic and works with any ODE integrator.
• Advantage‑Weighted Matching (AWM) – Weights the standard velocity‑matching loss by per‑sample advantages derived from reward signals, tightly coupling RL guidance with the original flow‑matching objective.

Experimental Validation
The authors evaluate Flow‑Factory on the large multimodal model Flux‑1‑dev (Labs, 2026) using three RL algorithms (Flow‑GRPO, DiffusionNFT, AWM) and two reward models (Text‑Rendering, PickScore).

• Reproduction of Published Results – Reward curves closely match those reported in the original papers, confirming that the decoupled architecture does not sacrifice stability or performance.
• Qualitative Improvements – Images generated after RL fine‑tuning display higher aesthetic quality and better alignment with human preferences compared to the base Flux model.
• Training Efficiency – With preprocessing‑based memory optimization on an 8‑GPU H200 cluster, peak memory drops from 61.08 GB to 53.14 GB per device, and per‑step time improves from 144 s to 82 s (1.74×). This enables larger batch sizes and faster iteration cycles on commodity hardware.

Conclusion and Outlook
Flow‑Factory delivers a practical solution to the engineering bottlenecks that have limited the broader adoption of RL for flow‑matching models. By cleanly separating algorithms, models, and rewards, providing a plug‑and‑play registry, and introducing memory‑saving preprocessing together with a robust multi‑reward loader, the framework allows researchers to prototype, benchmark, and scale new ideas with minimal code changes. The authors anticipate that the modular design will remain future‑proof as new foundation models (e.g., text‑to‑video, multimodal large language models) and novel RL paradigms (offline RL, meta‑RL) emerge, making Flow‑Factory a foundational infrastructure for the next generation of preference‑aligned generative AI.