Search papers, labs, and topics across Lattice.
This paper introduces LC-GRPO, a novel framework that addresses the train-inference gap in flow-based generative models by utilizing Langevin correction during policy optimization in reinforcement learning. By first applying an inference-aligned ODE Euler step followed by a stochastic Langevin correction, the method effectively aligns the distribution of training rollouts with that of deterministic inference, significantly improving the quality of generated samples. Experimental results across various tasks demonstrate that LC-GRPO enhances reward optimization while maintaining generation fidelity, outperforming traditional methods in bridging the gap between stochastic training and deterministic inference.
Langevin correction can transform the way flow-based generative models optimize policies, leading to sharper and more accurate sample generation during reinforcement learning.
Flow-based generative models are typically sampled by solving a deterministic ordinary differential equation (ODE), whereas online reinforcement learning requires stochastic rollouts for policy exploration and optimization. Existing GRPO methods for flow models therefore replace the inference-time ODE with a stochastic differential equation (SDE) during training. Although the ODE and SDE share the same marginal distributions in continuous time, their finite-step discretizations can differ substantially. In particular, SDE rollouts often become blurry as the exploration noise increases, creating a mismatch between the samples used for reinforcement learning and those generated by the test-time ODE sampler. We introduce LC-GRPO, a flow-based GRPO framework with Langevin correction. Each rollout transition first takes an inference-aligned ODE Euler step and then applies a stochastic Langevin correction targeting the marginal distribution at the resulting timestep. The required score is recovered directly from the flow velocity, requiring no additional score model, while the resulting transition remains an isotropic Gaussian with a tractable likelihood for policy optimization. We theoretically show that, under suitable conditions, one Langevin correction step reduces the Wasserstein error of an imperfect ODE Euler step. At a matched randomness level, we further show that the proposed transition can be more accurate than the standard Euler--Maruyama discretization of the reverse SDE. Experiments on SD3.5-Medium, FLUX.1-Dev, and HunyuanVideo demonstrate that LC-GRPO consistently improves reward optimization across text-to-image and text-to-video tasks, preserves generation quality, and substantially narrows the gap between stochastic training rollouts and deterministic test-time ODE inference.
