Flow matching stability improved by budgeting per-step variance in image generation
$λ$-Controlled GRPO: Turning Flow-Matching Ratio Instability into a Budgeted Resource
Machine Learning
Summary
Training AI models that create images based on rewards can become unstable during the step-by-step process of image generation. The authors identify that this instability comes from a single measurable factor called path variance, linked to the way the model samples data. Instead of trying to fix many symptoms separately, they control this path variance using a new method that predicts and budgets it throughout training. This approach helps the model better understand tricky text instructions and align with human preferences more reliably without relying on guesswork.
What this means in practice
- •For machine learning engineers: Train text-to-image generators with more stable reward-guided updates to improve text accuracy and human preference alignment.
- •For automated content moderation teams: Improve verification of generated text in images using stable reinforcement learning to better detect inaccuracies in optical character recognition tasks.
Authors
Yufeng Wang, Parivesh Priye, Meeshawn Marathe, Ramit Pahwa
Abstract
Reinforcement learning is increasingly used to align image generators with reward signals, and Flow-GRPO recently extended this paradigm to flow-matching models by treating the denoising sampler as a stochastic policy that can be optimized from reward feedback. Training in this setting is unstable in a way specific to multi-step denoising: the policy update changes systematically across denoising steps, with importance ratios drifting below one, becoming increasingly dispersed, clipping at different rates, and leaving fewer usable samples late in training. Prior work treats these effects as separate failure modes and addresses each with a hand-tuned stabilizer. We show instead that they arise from a single per-step quantity, which we call path variance. This quantity is determined exactly by the sampler's Gaussian transition kernel and can be estimated cheaply during training. This reframes instability as a resource that can be measured and budgeted rather than a collection of symptoms to repair. Our method, $λ$-Controlled GRPO, calibrates importance-ratio behavior from this predicted law rather than from noisy empirical statistics, and allocates gradient effort across denoising steps according to their predicted cost. The two scales governing the update are fixed by standard policy choices rather than introduced as free tuning parameters. On a text-to-image model under two reward settings, rendering difficult target text scored by optical character recognition and matching human preferences scored by a preference model, $λ$-Controlled GRPO improves both text accuracy and preference reward over the strongest empirical stabilizer. It also keeps late-step path variance within its intended budget, precisely where the baseline systematically overshoots. The result is a Flow-GRPO update calibrated by its own transition law rather than stabilized after instability appears.