NormGuard Preserves Image Quality in RL-Tuned Flow-Based Generators.

Reinforcement learning (RL) is often used to fine-tune flow-based generative models, improving their alignment with specific rewards. However, this post-training process frequently leads to a degradation in the perceptual quality of the generated outputs, which is not adequately captured by the reward proxy. A key finding is that RL fine-tuning tends to inflate the per-step velocity norm of the model, a structural signature of this quality drift. Existing inference-time corrections, such as rescaling the velocity norm, have been explored in other contexts but do not effectively transfer to RL. This is because the norm inflation becomes co-adapted into the model's weights, and simple rescaling at inference time neither improves the reward nor fixes the quality degradation. Furthermore, an analysis shows that velocity magnitude rescaling does not carry a consistent first-order reward signal at the batch level, suggesting that merely suppressing norm inflation won't inherently remove a reward-carrying component. These insights led to the development of NormGuard, a training-time intervention. NormGuard is a hinge penalty that activates only when the model's velocity norm exceeds a reference norm, and it is added to any velocity-local base loss. Experiments across various base models, post-training methods, and reward proxies demonstrate that NormGuard consistently enhances MLLM-judged image quality and forensic realism while fully preserving the intended reward. Its benefits are particularly amplified under few-step inference, and these improvements are not simply due to early stopping.