Discrete Flow Matching Policy Optimization

Maojiang Su; Po-Chung Hsieh; Weimin Wu; Mingcheng Lu; Jiunhau Chen; Jerry Yao-Chieh Hu; Han Liu

cs.LGApr 7, 2026

Discrete Flow Matching Policy Optimization

Authors: Maojiang Su, Po-Chung Hsieh, Weimin Wu, Mingcheng Lu, Jiunhau Chen, Jerry Yao-Chieh Hu, Han Liu

Organizations: Center for Foundation Models and Generative AI, Northwestern University, Evanston, IL 60208, USA · Department of Electrical Engineering, National Taiwan University, Taipei 10617, Taiwan · Department of Physics, National Taiwan University, Taipei 10617, Taiwan · Department of Statistics and Data Science, Northwestern University, Evanston, IL 60208, USA

Abstract

We introduce Discrete flow Matching policy Optimization (DoMinO), a unified framework for Reinforcement Learning (RL) fine-tuning Discrete Flow Matching (DFM) models under a broad class of policy gradient methods. Our key idea is to view the DFM sampling procedure as a multi-step Markov Decision Process. This perspective provides a simple and transparent reformulation of fine-tuning reward maximization as a robust RL objective. Consequently, it not only preserves the original DFM samplers but also avoids biased auxiliary estimators and likelihood surrogates used by many prior RL fine-tuning methods. To prevent policy collapse, we also introduce new total-variation regularizers to keep the fine-tuned distribution close to the pretrained one. Theoretically, we establish an upper bound on the discretization error of DoMinO and tractable upper bounds for the regularizers. Experimentally, we evaluate DoMinO on regulatory DNA sequence design. DoMinO achieves stronger predicted enhancer activity and better sequence naturalness than the previous best reward-driven baselines. The regularization further improves alignment with the natural sequence distribution while preserving strong functional performance. These results establish DoMinO as an useful framework for controllable discrete sequence generation.

Figures & tables

\displaystyle\theta\leftarrow\theta+\eta\frac{1}{M}\sum_{m=1}^{M}\sum_{t=0}^{T-\Delta t}\nabla_{\theta}\log p_{t}^{\theta}(x_{t+\Delta t}^{(m)}|x_{t}^{(m)})\widehat{A}^{(m)}

Algorithm 1 DoMinO-REINFORCE

\displaystyle\theta\leftarrow\theta+\eta\nabla_{\theta}\frac{1}{M}\sum_{m=1}^{M}\sum_{t=0}^{T-\Delta t}\min\left\{r_{t}^{(m)}(\theta)\widehat{A}_{t}^{(m)},\text{Clip}\left(r_{t}^{(m)}(\theta),1-\epsilon,1+\epsilon\right)\widehat{A}_{t}^{(m)}\right\}

Algorithm 2 DoMinO-PPO

Method	Pred-Activity $\uparrow$	ATAC-Acc $\uparrow$ (%)	3-mer Corr - All $\uparrow$
Pre-trained Diffusion	0.17 (0.01)	1.5 (0.3)	0.925 (0.004)
Pre-trained Flow-matching	0.64 (0.01)	1.1 (0.4)	0.884 (0.004)
DRAKES	6.37 (0.04)	96.1 (0.5)	-0.379 (0.009)
SEPO	7.55 (0.01)	99.5 (0.2)	-0.537 (0.002)
DoMinO-REINFORCE	8.32 (0.01)	99.2 (0.2)	-0.285 (0.001)
DoMinO-PPO	8.35 (0.00)	99.2 (0.2)	-0.331 (0.001)

Table 1: Performance of DNA design methods on the HepG2 enhancer design task. We compare pre-trained generative models, prior reward-driven baselines (DRAKES and SEPO), and our policy-gradient fine-tuning methods for discrete flow matching. We do not apply regularization loss during fine-tuning in this setting. Pred-Activity and ATAC-Acc evaluate functional performance, while 3-mer Corr-All measures sequence naturalness relative to the HepG2 data distribution. Our methods outperform the state-of-the-art baseline SEPO on Pred-Activity and 3-mer Corr-All and achieve comparable performance on ATAC-Acc. Bold and underlined entries highlight the best and second-best reward-driven methods.

Method	Pred-Activity $\uparrow$	ATAC-Acc $\uparrow$ (%)	3-mer Corr - All $\uparrow$
DRAKES	6.37 (0.04)	96.1 (0.5)	-0.379 (0.009)
DRAKES with KL	5.61 (0.07)	92.5 (0.6)	-0.302 (0.011)
SEPO	7.55 (0.01)	99.5 (0.2)	-0.537 (0.002)
SEPO with GF	7.64 (0.01)	99.9 (0.09)	-0.496 (0.001)
DoMinO-REINFORCE	8.32 (0.01)	99.2 (0.2)	-0.285 (0.001)
DoMinO-REINFORCE with CE	8.22 (0.01)	94.1 (0.8)	-0.347 (0.004)

Table 2: Effect of regularization on DoMinO for the HepG2 enhancer design task. We compare our policy-gradient fine-tuning methods with and without regularization loss. Pred-Activity and ATAC-Acc measure functional performance, while 3-mer Corr-All measures sequence naturalness relative to the HepG2 data distribution. Regularization improves sequence naturalness and helps maintain strong functional performance. This yields a better balance between reward optimization and naturalness. Bold and underlined entries highlight the best and second-best regularized reward-driven methods.

TV Regularization

See all 10 papers

dFlowGRPO: Rate-Aware Policy Optimization for Discrete Flow Models
May 10, 2026Zhengyan Wan, Yidong Ouyang, Panwen Hu +1Group Relative Policy Optimization Image Generation
Mohamed bin Zayed University of Artificial Intelligence · East China Normal University · University of California, Los Angeles +1
A Survey of Flow Matching in Reinforcement Learning
Oct 12, 2026Nabuat Zaman Nahim, Fairoz Nower Khan, Peizhong JuReinforcement Learning Flow-Based Generative Modeling
Department of Computer Science, University of Kentucky, Lexington, Kentucky 40506, USA
Discrete Flow Matching for Offline-to-Online Reinforcement Learning
May 12, 2026Fairoz Nower Khan, Nabuat Zaman Nahim, Peizhong JuFlow Matching Reinforcement Learning
Department of Computer Science, University of Kentucky

May 10, 2026J/K move · Enter open · S save

dFlowGRPO: Rate-Aware Policy Optimization for Discrete Flow Models

Zhengyan Wan, Yidong Ouyang, Panwen Hu, Qiang Sun

Mohamed bin Zayed University of Artificial Intelligence · East China Normal University · University of California, Los Angeles+1

Group Relative Policy Optimization Image Generation Flow-Based Generative Modeling RL for Generative Models Multimodal Understanding

Open paper AskPDF arXiv

Discrete flow models (DFMs) are a class of flexible generative models for generating discrete data, and diffusion large language models (dLLMs) can be viewed as a special case with a specific choice of a mixture path and a masked source distribution. While several recent works have explored reinforcement learning for dLLMs, its application to more general discrete flow models remains underexplored. In this work, we present discrete Flow-GRPO (dFlowGRPO), a unified reinforcement learning framework for discrete flow models that supports a broad family of probability paths and non-masked source distributions. We derive the full trajectory probability for DFMs and formulate the denoising process as a Markov decision process, enabling dFlowGRPO to incorporate information from both the associated conditional transition rates and the posterior model during reinforcement learning. We apply dFlowGRPO to FUDOKI, a recent multimodal discrete flow model, and evaluate it on both image generation and multimodal understanding tasks. Empirical results show that dFlowGRPO outperforms existing GRPO-type methods for dLLMs on text-to-image generation tasks and achieves performance competitive with continuous flow-based models trained using Flow-GRPO, while also demonstrating strong capabilities on understanding tasks.

Abstract

TV Regularization

Explore similar work

dFlowGRPO: Rate-Aware Policy Optimization for Discrete Flow Models

A Survey of Flow Matching in Reinforcement Learning

Discrete Flow Matching for Offline-to-Online Reinforcement Learning