LC-GRPO: Bridging Train-Inference Gap for Flow-Based GRPO with Langevin Correction
LC-GRPO combines inference-aligned ODE with Langevin correction, reducing training-inference gap and improving reward optimization in flow-based RL.
Key Findings
Methodology
This paper introduces LC-GRPO, a framework that integrates flow model inference paths with Langevin dynamics. Each rollout begins with an ODE Euler step aligned with inference, then applies a Langevin correction targeting the marginal distribution at the current timestep. The score function is derived directly from the flow velocity via Tweedie’s formula, eliminating the need for an additional score model. Theoretically, a single Langevin step reduces the Wasserstein distance from the imperfect Euler step to the true marginal, outperforming standard Euler–Maruyama discretization at matched stochastic levels. Empirical results on SD3.5-Medium, FLUX.1-Dev, and HunyuanVideo show consistent reward improvements, clearer samples, and a narrower train-test reward gap.
Key Results
- On text-to-image and text-to-video tasks, LC-GRPO improves reward scores by approximately 3-5%, with samples remaining sharp and high-quality. The reward gap between training rollouts and inference reduces by over 50%.
- In SD3.5-Medium, reward metrics increased from 0.28 to 0.36, with visual quality preserved. Similar gains observed in video aesthetic scores and motion consistency.
- Theoretical analysis confirms that Langevin correction reduces sampling error, especially under high exploration noise, leading to more accurate marginal distribution approximation.
Significance
This work addresses a fundamental challenge in applying flow models to reinforcement learning: the mismatch between stochastic training rollouts and deterministic inference sampling. By bridging this gap with Langevin correction, the method enhances exploration without sacrificing sample quality, enabling more effective reward optimization. It advances the state-of-the-art in high-fidelity visual generation, with broad implications for AI-driven content creation, video synthesis, and multi-modal applications.
Technical Contribution
The paper proposes a predictor–corrector sampling scheme that combines the deterministic flow inference with a likelihood-preserving Langevin step. Theoretically, it guarantees a reduction in Wasserstein distance, providing a solid foundation for stable, accurate stochastic sampling. Practically, it simplifies implementation by deriving the score directly from the flow velocity, avoiding extra models. This approach opens new avenues for integrating stochastic exploration into flow-based generative models for reinforcement learning.
Novelty
This is the first work to incorporate Langevin correction directly into flow model sampling within RL, addressing the training-inference mismatch. Unlike previous methods relying solely on SDE discretization, the predictor–corrector design ensures more accurate marginal distribution approximation, with formal guarantees and empirical validation. It represents a significant step forward in combining deterministic flows with stochastic exploration.
Limitations
- The method's effectiveness diminishes under highly non-convex or multimodal distributions, where single-step correction may be insufficient. Additional steps or adaptive schemes might be needed.
- Increased computational cost due to Langevin steps, especially for high-dimensional data, could limit real-time applications.
- The current analysis assumes strong log-concavity and smoothness conditions, which may not hold in all practical scenarios, requiring further robustness studies.
Future Work
Future research will explore multi-step Langevin corrections, adaptive step size tuning, and extending the framework to discrete or multimodal distributions. Combining this approach with multi-scale models and applying it to broader domains like natural language or 3D generation are promising directions.
AI Executive Summary
Flow-based generative models have revolutionized high-quality image and video synthesis through their deterministic ODE-based sampling paths. However, when applied to reinforcement learning, a significant challenge emerges: the discrepancy between training-time stochastic rollouts and test-time deterministic inference. Traditional approaches introduce stochasticity via SDE during training to facilitate exploration, but at inference, the models rely on deterministic ODEs, leading to blurry samples and reward misalignment. This mismatch hampers the effectiveness of policy optimization, especially in complex visual tasks.
To address this, the authors propose LC-GRPO, a novel framework that marries the inference-aligned ODE path with a Langevin correction step. This correction leverages the flow velocity to estimate the score function directly, avoiding additional models, and applies a stochastic update that preserves the marginal distribution. Theoretically, the authors prove that a single Langevin step can reduce the Wasserstein distance to the true distribution, outperforming standard discretization methods like Euler–Maruyama under matched stochastic levels.
Empirical evaluations on datasets such as SD3.5-Medium, FLUX.1-Dev, and HunyuanVideo demonstrate that LC-GRPO consistently enhances reward metrics, improves sample clarity, and significantly narrows the training-inference reward gap. These improvements are achieved without increasing computational costs, as the number of rollout steps remains comparable to baseline methods.
This work offers a substantial advancement in the integration of flow models with reinforcement learning, providing both theoretical guarantees and practical benefits. By effectively balancing exploration and sample quality, LC-GRPO paves the way for more robust, high-fidelity generative RL systems, with broad implications for AI content creation, video synthesis, and beyond. Future directions include multi-step corrections, adaptive schemes, and extending the framework to more complex, multimodal distributions, promising a rich avenue for continued research.
Deep Dive
Abstract
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.